A lot of stuff

RupertAvery · Aug 26, 2024 · 6fa0e64 · 6fa0e64
1 parent f65b7d4
commit 6fa0e64
Show file tree

Hide file tree

Showing 36 changed files with 1,036 additions and 602 deletions.
diff --git a/Diffusion.Database/DataStore.cs b/Diffusion.Database/DataStore.cs
@@ -85,6 +85,11 @@ await Task.Run(() =>
         db.CreateIndex<Image>(image => image.FileSize);
         db.CreateIndex<Image>(image => image.ModifiedDate);
 
+        db.CreateIndex<Image>(image => image.Prompt);
+        db.CreateIndex<Image>(image => image.NegativePrompt);
+        db.CreateIndex<Image>(image => image.Workflow);
+        db.CreateIndex<Image>(image => image.WorkflowId);
+
         db.CreateIndex("Image", new[] { "ForDeletion", "CreatedDate" });
         db.CreateIndex("Image", new[] { "NSFW", "CreatedDate" });
         db.CreateIndex("Image", new[] { "Unavailable", "CreatedDate" });

diff --git a/Diffusion.Database/Image.cs b/Diffusion.Database/Image.cs
@@ -43,6 +43,8 @@ public class Image
     public int? ENSD { get; set; }
     public long FileSize { get; set; }
     public bool NoMetadata { get; set; }
+    public string? Workflow { get; set; }
+    public string? WorkflowId { get; set; }
 }
 
 

diff --git a/Diffusion.Database/QueryBuilder.Filter.cs b/Diffusion.Database/QueryBuilder.Filter.cs
@@ -62,7 +62,6 @@ public static (string WhereClause, IEnumerable<object> Bindings, IEnumerable<str
             FilterNegativePromptEx(filter, conditions);
             FilterPromptEx(filter, conditions);
 
-
             return (string.Join(" AND ", conditions.Select(c => c.Key)),
                 conditions.SelectMany(c =>
                 {

diff --git a/Diffusion.Scanner/ComfyUI.cs b/Diffusion.Scanner/ComfyUI.cs
@@ -0,0 +1,105 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Text.Json;
+using System.Threading.Tasks;
+using System.Xml.Linq;
+
+namespace Diffusion.IO
+{
+    public class Node
+    {
+        public string Id { get; set; }
+        public string Name { get; set; }
+        public List<Input> Inputs { get; set; }
+
+        public override int GetHashCode()
+        {
+            var hash = Id.GetHashCode();
+
+            hash = (hash * 397) ^ Name.GetHashCode();
+
+            foreach (var input in Inputs)
+            {
+                hash = (hash * 397) ^ input.Name.GetHashCode();
+            }
+
+            return hash;
+        }
+    }
+
+    public class Input
+    {
+        public string Name { get; set; }
+        public object Value { get; set; }
+
+        public Input(string name, object value)
+        {
+            Name = name;
+            Value = value;
+        }
+    }
+
+    public class ComfyUIParser
+    {
+        public IReadOnlyCollection<Node> Parse(string workflow)
+        {
+            if (workflow == null) return null;
+            var root = System.Text.Json.JsonDocument.Parse(workflow);
+            var rootElements = root.RootElement.EnumerateObject().ToDictionary(n => n.Name, n => n.Value);
+
+            var nodes = new List<Node>();
+
+            foreach (var element in rootElements)
+            {
+                var node = new Node();
+                node.Id = element.Key;
+
+                foreach (var props in element.Value.EnumerateObject())
+                {
+
+                    if (props.Name == "inputs")
+                    {
+                        node.Inputs = new List<Input>();
+                        foreach (var prop2 in props.Value.EnumerateObject())
+                        {
+                            switch (prop2.Value.ValueKind)
+                            {
+                                case JsonValueKind.Undefined:
+                                    break;
+                                case JsonValueKind.Object:
+                                    break;
+                                case JsonValueKind.Array:
+                                    break;
+                                case JsonValueKind.String:
+                                    node.Inputs.Add(new Input(prop2.Name.Replace("_", "__"), prop2.Value.GetString()));
+                                    break;
+                                case JsonValueKind.Number:
+                                    node.Inputs.Add(new Input(prop2.Name.Replace("_", "__"), prop2.Value.GetDouble()));
+                                    break;
+                                case JsonValueKind.True:
+                                    node.Inputs.Add(new Input(prop2.Name.Replace("_", "__"), prop2.Value.GetBoolean()));
+                                    break;
+                                case JsonValueKind.False:
+                                    node.Inputs.Add(new Input(prop2.Name.Replace("_", "__"), prop2.Value.GetBoolean()));
+                                    break;
+                                case JsonValueKind.Null:
+                                    break;
+                                default:
+                                    throw new ArgumentOutOfRangeException();
+                            }
+                        }
+                    }
+                    else if (props.Name == "class_type")
+                    {
+                        node.Name = props.Value.GetString().Replace("_", "__");
+                    }
+                }
+                nodes.Add(node);
+            }
+
+            return nodes;
+        }
+    }
+}
diff --git a/Diffusion.Scanner/FileParameters.cs b/Diffusion.Scanner/FileParameters.cs
@@ -25,4 +25,6 @@ public class FileParameters
     public decimal? PromptStrength { get; set; }
     public long FileSize { get; set; }
     public bool NoMetadata { get; set; }
+    public string? Workflow { get; set; }
+    public string? WorkflowId { get; set; }
 }