Fix broken NovelAI metadata reading introduced by FooocusMRE.

#175
RupertAvery · Jan 12, 2024 · 51ce55f · 51ce55f
1 parent f8ed439
commit 51ce55f
Showing 1 changed file with 37 additions and 13 deletions.
diff --git a/Diffusion.Scanner/Metadata.cs b/Diffusion.Scanner/Metadata.cs
@@ -8,6 +8,7 @@
 using System.Globalization;
 using MetadataExtractor.Formats.WebP;
 using Diffusion.Common;
+using System.Linq;
 
 namespace Diffusion.IO;
 
@@ -57,6 +58,8 @@ public enum MetaFormat
 
                     decimal aestheticScore = 0;
 
+                    //string tagData = null;
+
                     foreach (var directory in directories)
                     {
                         if (directory.Name == "PNG-tEXt")
@@ -67,19 +70,36 @@ public enum MetaFormat
                                 {
                                     if (tag.Description.StartsWith("parameters:"))
                                     {
-                                        var isJson = tag.Description.Substring("parameters: ".Length).Trim().StartsWith("{");
-                                        format = isJson ? MetaFormat.RuinedFooocus : MetaFormat.A1111;
-                                        fileParameters = isJson ? ReadRuinedFooocusParameters(tag.Description) : ReadA111Parameters(tag.Description);
+                                        if (fileParameters == null)
+                                        {
+                                            var isJson = tag.Description.Substring("parameters: ".Length).Trim().StartsWith("{");
+                                            format = isJson ? MetaFormat.RuinedFooocus : MetaFormat.A1111;
+                                            fileParameters = isJson ? ReadRuinedFooocusParameters(tag.Description) : ReadA111Parameters(tag.Description);
+                                        }
                                     }
                                     else if (tag.Description.StartsWith("Comment:"))
                                     {
-                                        format = MetaFormat.FooocusMRE;
-                                        fileParameters = ReadFooocusMREParameters(tag.Description);
+                                        if (fileParameters == null)
+                                        {
+                                            if (directory.Tags.Any(t => t.Description == "Software: NovelAI"))
+                                            {
+                                                format = MetaFormat.NovelAI;
+                                                fileParameters = ReadNovelAIParameters(file, directories);
+                                            }
+                                            else
+                                            {
+                                                format = MetaFormat.FooocusMRE;
+                                                fileParameters = ReadFooocusMREParameters(tag.Description);
+                                            }
+                                        }
                                     }
                                     else if (tag.Description == "Software: NovelAI")
                                     {
-                                        format = MetaFormat.NovelAI;
-                                        fileParameters = ReadNovelAIParameters(file, directories);
+                                        if (fileParameters == null)
+                                        {
+                                            format = MetaFormat.NovelAI;
+                                            fileParameters = ReadNovelAIParameters(file, directories);
+                                        }
                                     }
                                     else if (tag.Description.StartsWith("Dream: "))
                                     {
@@ -98,9 +118,12 @@ public enum MetaFormat
                                     }
                                     else if (tag.Description.StartsWith("prompt: "))
                                     {
-                                        var isJson = tag.Description.Substring("prompt: ".Length).Trim().StartsWith("{");
-                                        format = isJson ? MetaFormat.ComfyUI : MetaFormat.EasyDiffusion;
-                                        fileParameters = isJson ? ReadComfyUIParameters(file, tag.Description) : ReadEasyDiffusionParameters(file, directories);
+                                        if (fileParameters == null)
+                                        {
+                                            var isJson = tag.Description.Substring("prompt: ".Length).Trim().StartsWith("{");
+                                            format = isJson ? MetaFormat.ComfyUI : MetaFormat.EasyDiffusion;
+                                            fileParameters = isJson ? ReadComfyUIParameters(file, tag.Description) : ReadEasyDiffusionParameters(file, directories);
+                                        }
                                     }
                                     else if (tag.Description.StartsWith("Score:"))
                                     {
@@ -153,7 +176,7 @@ public enum MetaFormat
                     {
                         fileParameters ??= new FileParameters();
                         fileParameters.AestheticScore = aestheticScore;
-                        if(fileParameters.OtherParameters == null)
+                        if (fileParameters.OtherParameters == null)
                         {
                             fileParameters.OtherParameters = $"aesthetic_score: {fileParameters.AestheticScore}";
                         }
@@ -653,6 +676,7 @@ private static FileParameters ReadNovelAIParameters(string file, IEnumerable<Dir
 
         var propList = new List<string>();
 
+
         if (TryFindTag(directories, "PNG-tEXt", "Textual Data", tag => tag.Description.StartsWith("Comment:"), out tag))
         {
             var json = JsonDocument.Parse(tag.Description.Substring("Comment: ".Length));
@@ -666,7 +690,7 @@ private static FileParameters ReadNovelAIParameters(string file, IEnumerable<Dir
 
             foreach (var property in properties)
             {
-                if (property.Name != "uc")
+                if (property.Name != "uc" && property.Name != "prompt")
                 {
                     propList.Add($"{property.Name}: {property.Value.ToString()}");
                 }
@@ -1105,7 +1129,7 @@ private static FileParameters ReadAutomatic1111Parameters(string file, IEnumerab
         {
             fileParameters ??= new FileParameters();
             fileParameters.AestheticScore = aestheticScore;
-            if(fileParameters.OtherParameters == null)
+            if (fileParameters.OtherParameters == null)
             {
                 fileParameters.OtherParameters = $"aesthetic_score: {fileParameters.AestheticScore}";
             }