Rework union type in workflow schema

the backendtype being simple settable from outside was required in settingsImageWorkflowSelector. Signed-off-by: Florian Esser <florian.esser@tngtech.com>
TNG · Dec 4, 2024 · b87b32e · b87b32e
1 parent a962981
commit b87b32e
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 49 deletions.
diff --git a/WebUI/external/workflows/fluxQ4.json b/WebUI/external/workflows/fluxQ4.json
@@ -4,18 +4,17 @@
         "Q4",
         "Fast"
     ],
-    "backend": {
-        "comfyUI": {
-            "customNodes": [
-                "city96/ComfyUI-GGUF"
-            ],
-            "requiredModels": [
-                "unet:city96/FLUX.1-schnell-gguf/flux1-schnell-Q4_K_S.gguf",
-                "clip:city96/t5-v1_1-xxl-encoder-gguf/t5-v1_1-xxl-encoder-Q3_K_M.gguf",
-                "clip:comfyanonymous/flux_text_encoders/clip_l.safetensors",
-                "vae:black-forest-labs/FLUX.1-schnell/ae.safetensors"
-            ]
-        }
+    "backend": "comfyui",
+    "comfyUIRequirements": {
+        "customNodes": [
+            "city96/ComfyUI-GGUF"
+        ],
+        "requiredModels": [
+            "unet:city96/FLUX.1-schnell-gguf/flux1-schnell-Q4_K_S.gguf",
+            "clip:city96/t5-v1_1-xxl-encoder-gguf/t5-v1_1-xxl-encoder-Q3_K_M.gguf",
+            "clip:comfyanonymous/flux_text_encoders/clip_l.safetensors",
+            "vae:black-forest-labs/FLUX.1-schnell/ae.safetensors"
+        ]
     },
 
     "requirements": [

diff --git a/WebUI/external/workflows/fluxQ8.json b/WebUI/external/workflows/fluxQ8.json
@@ -1,20 +1,18 @@
 {
     "name": "Flux.1-Schnell High Quality",
 
-    "backend": {
-        "comfyUI": {
-            "customNodes": [
-                "city96/ComfyUI-GGUF"
-            ],
-            "requiredModels": [
-                "unet:city96/FLUX.1-schnell-gguf/flux1-schnell-Q8_0.gguf",
-                "clip:city96/t5-v1_1-xxl-encoder-gguf/t5-v1_1-xxl-encoder-Q3_K_M.gguf",
-                "clip:comfyanonymous/flux_text_encoders/clip_l.safetensors",
-                "vae:black-forest-labs/FLUX.1-schnell/ae.safetensors"
-            ]
-        }
+    "backend": "comfyui",
+    "comfyUIRequirements": {
+        "customNodes": [
+            "city96/ComfyUI-GGUF"
+        ],
+        "requiredModels": [
+            "unet:city96/FLUX.1-schnell-gguf/flux1-schnell-Q8_0.gguf",
+            "clip:city96/t5-v1_1-xxl-encoder-gguf/t5-v1_1-xxl-encoder-Q3_K_M.gguf",
+            "clip:comfyanonymous/flux_text_encoders/clip_l.safetensors",
+            "vae:black-forest-labs/FLUX.1-schnell/ae.safetensors"
+        ]
     },
-
     "tags": [
         "Q8",
         "Fast"

diff --git a/WebUI/src/assets/js/store/imageGeneration.ts b/WebUI/src/assets/js/store/imageGeneration.ts
@@ -50,12 +50,6 @@ export type Setting = z.infer<typeof SettingSchema>
 
 const WorkflowRequirementSchema = z.enum(['high-vram'])
 
-const comfyUIBackendSchema = z.object({comfyUI: z.object({
-        customNodes: z.array(z.string()),
-        requiredModels: z.array(z.string()),
-    })})
-export type comfyUIBackend = z.infer<typeof comfyUIBackendSchema>
-
 const ComfyUIApiWorkflowSchema = z.record(z.string(), z.object({
     inputs: z.object({
         text: z.string().optional(),
@@ -67,10 +61,11 @@ export type ComfyUIApiWorkflow = z.infer<typeof ComfyUIApiWorkflowSchema>;
 
 const WorkflowSchema = z.object({
     name: z.string(),
-    backend: z.union([
-        comfyUIBackendSchema,
-        z.enum(['default']),
-        ]),
+    backend: z.enum(['default', 'comfyui']),
+    comfyUIRequirements: z.object({
+        customNodes: z.array(z.string()),
+        requiredModels: z.array(z.string()),
+    }).optional(),
     tags: z.array(z.string()),
     requiredModels: z.array(z.string()).optional(),
     requirements: z.array(WorkflowRequirementSchema),
@@ -90,13 +85,6 @@ const WorkflowSchema = z.object({
 })
 export type Workflow = z.infer<typeof WorkflowSchema>;
 
-function isDefaultBackend(data: Workflow): boolean {
-    return data.backend === "default"
-}
-
-function isComfyUiBackend(data: string | comfyUIBackend): data is comfyUIBackend {
-    return data !== "default"
-}
 
 const globalDefaultSettings = {
     width: 512,
@@ -497,16 +485,15 @@ export const useImageGeneration = defineStore("imageGeneration", () => {
     }
 
     async function getMissingModels() {
-        if (isComfyUiBackend(activeWorkflow.value!.backend)) {
-            return getMissingComfyuiBackendModels(activeWorkflow.value!.backend)
-        }
-        else {
+        if (activeWorkflow.value!.backend === "default") {
             return getMissingDefaultBackendModels()
+        } else {
+            return getMissingComfyuiBackendModels()
         }
     }
 
-    async function getMissingComfyuiBackendModels(comfyUIBackend: comfyUIBackend) {
-        if (comfyUIBackend.comfyUI.requiredModels === undefined) {
+    async function getMissingComfyuiBackendModels() {
+        if (activeWorkflow.value!.comfyUIRequirements!.requiredModels === undefined) {
             toast.error('Defined workflow did not specify required models. Please add "requiredModels" to workflowfile.');
             return []
         } else {
@@ -524,7 +511,7 @@ export const useImageGeneration = defineStore("imageGeneration", () => {
                 }
                 return {type: modelTypeToId(modelType), repo_id: repoAddress, backend: "comfyui"}
             }
-            const checkList: CheckModelAlreadyLoadedParameters[] = comfyUIBackend.comfyUI.requiredModels.map( extractDownloadModelParamsFromString)
+            const checkList: CheckModelAlreadyLoadedParameters[] = activeWorkflow.value!.comfyUIRequirements!.requiredModels.map( extractDownloadModelParamsFromString)
             const checkedModels: CheckModelAlreadyLoadedResult[]  = await globalSetup.checkModelAlreadyLoaded(checkList);
             const modelsToBeLoaded = checkedModels.filter(checkModelExistsResult => !checkModelExistsResult.already_loaded)
             for (const item of modelsToBeLoaded) {