Fix workflows + add support for hidream

Author: mythz

AiServer.ServiceInterface/GenerationServices.cs

@@ -298,10 +298,10 @@ public static async Task<GenerationResponse> ProcessSyncGenerationAsync(this Cre
 
         var completedResponse = new GenerationResponse { };
 
-        // Wait for the job to complete max 1 minute
-        var timeout = DateTime.UtcNow.AddMinutes(1);
+        // Wait for the job to complete max 20 minutes
+        var timeout = DateTime.UtcNow.AddSeconds(job.Job?.TimeoutSecs ?? 20 * 60);
         while (queuedJob?.Job?.State is not (BackgroundJobState.Completed or BackgroundJobState.Cancelled
-               or BackgroundJobState.Failed) && DateTime.UtcNow < timeout)
+                   or BackgroundJobState.Failed) && DateTime.UtcNow < timeout)
         {
             await Task.Delay(1000);
             queuedJob = jobs.GetJob(diffResponse.Id);

AiServer.ServiceInterface/VideoServices.cs

@@ -374,7 +374,7 @@ public static async Task<ArtifactGenerationResponse> ProcessSyncTransformAsync(t
 
         var completedResponse = new ArtifactGenerationResponse();
 
-        // Wait for the job to complete max 1 minute
+        // Wait for the job to complete max 20 minutes
         var timeout = DateTime.UtcNow.AddSeconds(job.Job?.TimeoutSecs ?? 20 * 60);
         while (queuedJob?.Job?.State is not (
                    BackgroundJobState.Completed or BackgroundJobState.Cancelled or BackgroundJobState.Failed) 

AiServer.ServiceModel/Types/Generations.cs

@@ -144,6 +144,8 @@ public class MediaModel : IHasId<string>
     public int? Steps { get; set; }
     public string? NegativePrompt { get; set; }
     public ModelType? ModelType { get; set; }
+    public string? Filename { get; set; }
+    public string? Workflow { get; set; }
     /// <summary>
     /// Flag for models that are pulled down when run for the first time
     /// Keys are the ProviderId strings

AiServer/Configure.AppHost.cs

@@ -57,14 +57,8 @@ public void Configure(IWebHostBuilder builder) => builder
             services.AddSingleton<CustomOpenAiProvider>();
             services.AddSingleton<AiProviderFactory>();
 
-            services.AddSingleton(new ComfyMediaProviderOptions
-            {
-                TextToImageModelOverrides = new Dictionary<string, string>
-                {
-                    ["flux1-schnell.safetensors"] = "flux1/text_to_image.json",
-                    ["sd3.5_large_fp8_scaled.safetensors"] = "sd35/text_to_image.json"
-                }
-            });
+            services.AddSingleton(new ComfyMediaProviderOptions());
+            
             services.AddSingleton<ComfyProvider>();
             services.AddSingleton<ReplicateAiProvider>();
             services.AddSingleton<OpenAiMediaProvider>();

AiServer/Configure.BackgroundJobs.cs

@@ -1,4 +1,5 @@
 using AiServer.ServiceInterface;
+using AiServer.ServiceInterface.Generation;
 using AiServer.ServiceInterface.Recurring;
 using ServiceStack.Data;
 using ServiceStack.Jobs;
@@ -32,6 +33,12 @@ public void Configure(IWebHostBuilder builder) => builder
             using var db = services.GetRequiredService<IDbConnectionFactory>().OpenDbConnection();
             AppData.Instance.Reload(db);
 
+            var mediaOptions = services.GetRequiredService<ComfyMediaProviderOptions>();
+            foreach (var model in AppData.Instance.MediaModels.Where(x => x.Filename != null && x.Workflow != null))
+            {
+                mediaOptions.TextToImageModelOverrides[model.Filename!] = model.Workflow!;
+            }
+            
             var jobs = services.GetRequiredService<IBackgroundJobs>();
             // #if DEBUG
             // jobs.RecurringCommand<LogCommand>("Every Minute", Schedule.EveryMinute, new() {

AiServer/workflows/hidream/text_to_image.json

@@ -97,30 +97,6 @@
       },
       "widgets_values": []
     },
-    {
-      "id": 72,
-      "type": "MarkdownNote",
-      "pos": [
-        1250,
-        -130
-      ],
-      "size": [
-        320,
-        360
-      ],
-      "flags": {},
-      "order": 1,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [],
-      "title": "Sampling Settings",
-      "properties": {},
-      "widgets_values": [
-        "## Official sampling settings\n\nProvided for reference, my workflows may have slightly different settings.\n\n### HiDream Full\n\n* hidream_i1_full_fp16.safetensors\n* shift: 3.0\n* steps: 50\n* sampler: uni_pc\n* scheduler: simple\n* cfg: 5.0\n\n### HiDream Dev\n\n* hidream_i1_dev_bf16.safetensors\n* shift: 6.0\n* steps: 28\n* sampler: lcm\n* scheduler: normal\n* cfg: 1.0 (no negative prompt)\n\n### HiDream Fast\n\n* hidream_i1_fast_bf16.safetensors\n* shift: 3.0\n* steps: 16\n* sampler: lcm\n* scheduler: normal\n* cfg: 1.0 (no negative prompt)\n"
-      ],
-      "color": "#432",
-      "bgcolor": "#653"
-    },
     {
       "id": 40,
       "type": "CLIPTextEncode",
@@ -211,75 +187,6 @@
         "ComfyUI"
       ]
     },
-    {
-      "id": 79,
-      "type": "MarkdownNote",
-      "pos": [
-        -100,
-        -110
-      ],
-      "size": [
-        240,
-        120
-      ],
-      "flags": {},
-      "order": 2,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [],
-      "properties": {},
-      "widgets_values": [
-        "You can try changing the weight_dtype to fp8 if you are running out of memory.\n\nDownload one of diffusion Models file [here](https://huggingface.co/Comfy-Org/HiDream-I1_ComfyUI/tree/main/split_files/diffusion_models)\n\nPut it in the **ComfyUI/models/diffusion_models** directory"
-      ],
-      "color": "#432",
-      "bgcolor": "#653"
-    },
-    {
-      "id": 78,
-      "type": "MarkdownNote",
-      "pos": [
-        -100,
-        50
-      ],
-      "size": [
-        220,
-        120
-      ],
-      "flags": {},
-      "order": 3,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [],
-      "properties": {},
-      "widgets_values": [
-        "Download all of the text_encoders models [here](https://huggingface.co/Comfy-Org/HiDream-I1_ComfyUI/tree/main/split_files/text_encoders)\n\nPut them in the **ComfyUI/models/clip** directory"
-      ],
-      "color": "#432",
-      "bgcolor": "#653"
-    },
-    {
-      "id": 80,
-      "type": "MarkdownNote",
-      "pos": [
-        -100,
-        230
-      ],
-      "size": [
-        210,
-        88
-      ],
-      "flags": {},
-      "order": 4,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [],
-      "properties": {},
-      "widgets_values": [
-        "Download the VAE model [here](https://huggingface.co/Comfy-Org/HiDream-I1_ComfyUI/blob/main/split_files/vae/ae.safetensors) \n\nPut it in the **ComfyUI/models/vae** directory"
-      ],
-      "color": "#432",
-      "bgcolor": "#653"
-    },
     {
       "id": 70,
       "type": "ModelSamplingSD3",

AiServer/wwwroot/lib/data/media-models.json

@@ -43,7 +43,6 @@
     "sampler": "dpmpp_2m_sde_gpu",
     "width": 1024,
     "height": 1024,
-    "type": "sd3",
     "url": "https://civitai.com/models/194768?modelVersionId=1031794",
     "downloadUrl": "https://civitai.com/api/download/models/1031794?type=Model&format=SafeTensor&size=pruned&fp=fp16"
   },
@@ -68,7 +67,7 @@
     "sampler": "euler",
     "width": 1024,
     "height": 1024,
-    "type": "flux1",
+    "workflow": "flux1/text_to_image.json",
     "downloadApiKeyVar": "HF_TOKEN",
     "url": "https://huggingface.co/black-forest-labs/FLUX.1-schnell",
     "downloadUrl": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors?download=true",
@@ -94,7 +93,7 @@
     "sampler": "lcm",
     "width": 1024,
     "height": 1024,
-    "type": "hidream",
+    "workflow": "hidream/text_to_image.json",
     "url": "https://comfyanonymous.github.io/ComfyUI_examples/hidream/",
     "downloadUrl": "https://huggingface.co/Comfy-Org/HiDream-I1_ComfyUI/resolve/main/split_files/diffusion_models/hidream_i1_dev_fp8.safetensors?download=true"
   },
@@ -116,7 +115,6 @@
     "sampler": "dpmpp_2m",
     "width": 1024,
     "height": 1024,
-    "type": "sd3",
     "url": "https://huggingface.co/stabilityai/stable-diffusion-3.5-large",
     "downloadUrl": "https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/sd3.5_large.safetensors?download=true"
   },
@@ -138,7 +136,7 @@
     "sampler": "dpmpp_2m",
     "width": 1024,
     "height": 1024,
-    "type": "sd3",
+    "workflow": "sd35/text_to_image.json",
     "url": "https://huggingface.co/stabilityai/stable-diffusion-3.5-large-turbo",
     "downloadUrl": "https://huggingface.co/stabilityai/stable-diffusion-3.5-large-turbo/resolve/main/sd3.5_large_turbo.safetensors?download=true"
   },
@@ -161,7 +159,6 @@
     "sampler": "euler",
     "width": 1024,
     "height": 1024,
-    "type": "sd3",
     "url": "https://huggingface.co/Comfy-Org/stable-diffusion-3.5-fp8",
     "downloadUrl": "https://huggingface.co/Comfy-Org/stable-diffusion-3.5-fp8/resolve/main/sd3.5_large_fp8_scaled.safetensors?download=true"
   },