Jelajahi Sumber

refac: image gen

Timothy J. Baek 8 bulan lalu
induk
melakukan
95057d2368

+ 144 - 195
backend/apps/images/main.py

@@ -1,5 +1,3 @@
-import re
-import requests
 from fastapi import (
 from fastapi import (
     FastAPI,
     FastAPI,
     Request,
     Request,
@@ -7,14 +5,6 @@ from fastapi import (
     HTTPException,
     HTTPException,
 )
 )
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.cors import CORSMiddleware
-
-from constants import ERROR_MESSAGES
-from utils.utils import (
-    get_verified_user,
-    get_admin_user,
-)
-
-from apps.images.utils.comfyui import ComfyUIGenerateImageForm, comfyui_generate_image
 from typing import Optional
 from typing import Optional
 from pydantic import BaseModel
 from pydantic import BaseModel
 from pathlib import Path
 from pathlib import Path
@@ -23,7 +13,21 @@ import uuid
 import base64
 import base64
 import json
 import json
 import logging
 import logging
+import re
+import requests
+
+from utils.utils import (
+    get_verified_user,
+    get_admin_user,
+)
+
+from apps.images.utils.comfyui import (
+    ComfyUIWorkflow,
+    ComfyUIGenerateImageForm,
+    comfyui_generate_image,
+)
 
 
+from constants import ERROR_MESSAGES
 from config import (
 from config import (
     SRC_LOG_LEVELS,
     SRC_LOG_LEVELS,
     CACHE_DIR,
     CACHE_DIR,
@@ -76,176 +80,178 @@ app.state.config.IMAGE_SIZE = IMAGE_SIZE
 app.state.config.IMAGE_STEPS = IMAGE_STEPS
 app.state.config.IMAGE_STEPS = IMAGE_STEPS
 
 
 
 
-def get_automatic1111_api_auth():
-    if app.state.config.AUTOMATIC1111_API_AUTH is None:
-        return ""
-    else:
-        auth1111_byte_string = app.state.config.AUTOMATIC1111_API_AUTH.encode("utf-8")
-        auth1111_base64_encoded_bytes = base64.b64encode(auth1111_byte_string)
-        auth1111_base64_encoded_string = auth1111_base64_encoded_bytes.decode("utf-8")
-        return f"Basic {auth1111_base64_encoded_string}"
-
-
 @app.get("/config")
 @app.get("/config")
 async def get_config(request: Request, user=Depends(get_admin_user)):
 async def get_config(request: Request, user=Depends(get_admin_user)):
     return {
     return {
-        "engine": app.state.config.ENGINE,
         "enabled": app.state.config.ENABLED,
         "enabled": app.state.config.ENABLED,
+        "engine": app.state.config.ENGINE,
+        "openai": {
+            "OPENAI_API_BASE_URL": app.state.config.OPENAI_API_BASE_URL,
+            "OPENAI_API_KEY": app.state.config.OPENAI_API_KEY,
+        },
+        "automatic1111": {
+            "AUTOMATIC1111_BASE_URL": app.state.config.AUTOMATIC1111_BASE_URL,
+            "AUTOMATIC1111_API_AUTH": app.state.config.AUTOMATIC1111_API_AUTH,
+        },
+        "comfyui": {
+            "COMFYUI_BASE_URL": app.state.config.COMFYUI_BASE_URL,
+            "COMFYUI_WORKFLOW": app.state.config.COMFYUI_WORKFLOW,
+            "COMFYUI_WORKFLOW_NODES": app.state.config.COMFYUI_WORKFLOW_NODES,
+        },
     }
     }
 
 
 
 
-class ConfigUpdateForm(BaseModel):
-    engine: str
+class OpenAIConfigForm(BaseModel):
+    OPENAI_API_BASE_URL: str
+    OPENAI_API_KEY: str
+
+
+class Automatic1111ConfigForm(BaseModel):
+    AUTOMATIC1111_BASE_URL: str
+    AUTOMATIC1111_API_AUTH: str
+
+
+class ComfyUIConfigForm(BaseModel):
+    COMFYUI_BASE_URL: str
+    COMFYUI_WORKFLOW: str
+    COMFYUI_WORKFLOW_NODES: list[dict]
+
+
+class ConfigForm(BaseModel):
     enabled: bool
     enabled: bool
+    engine: str
+    openai: OpenAIConfigForm
+    automatic1111: Automatic1111ConfigForm
+    comfyui: ComfyUIConfigForm
 
 
 
 
 @app.post("/config/update")
 @app.post("/config/update")
-async def update_config(form_data: ConfigUpdateForm, user=Depends(get_admin_user)):
+async def update_config(form_data: ConfigForm, user=Depends(get_admin_user)):
     app.state.config.ENGINE = form_data.engine
     app.state.config.ENGINE = form_data.engine
     app.state.config.ENABLED = form_data.enabled
     app.state.config.ENABLED = form_data.enabled
-    return {
-        "engine": app.state.config.ENGINE,
-        "enabled": app.state.config.ENABLED,
-    }
 
 
+    app.state.config.OPENAI_API_BASE_URL = form_data.openai.OPENAI_API_BASE_URL
+    app.state.config.OPENAI_API_KEY = form_data.openai.OPENAI_API_KEY
 
 
-class EngineUrlUpdateForm(BaseModel):
-    AUTOMATIC1111_BASE_URL: Optional[str] = None
-    AUTOMATIC1111_API_AUTH: Optional[str] = None
-    COMFYUI_BASE_URL: Optional[str] = None
+    app.state.config.AUTOMATIC1111_BASE_URL = (
+        form_data.automatic1111.AUTOMATIC1111_BASE_URL
+    )
+    app.state.config.AUTOMATIC1111_API_AUTH = (
+        form_data.automatic1111.AUTOMATIC1111_API_AUTH
+    )
 
 
+    app.state.config.COMFYUI_BASE_URL = form_data.comfyui.COMFYUI_BASE_URL
+    app.state.config.COMFYUI_WORKFLOW = form_data.comfyui.COMFYUI_WORKFLOW
+    app.state.config.COMFYUI_WORKFLOW_NODES = form_data.comfyui.COMFYUI_WORKFLOW_NODES
 
 
-@app.get("/url")
-async def get_engine_url(user=Depends(get_admin_user)):
     return {
     return {
-        "AUTOMATIC1111_BASE_URL": app.state.config.AUTOMATIC1111_BASE_URL,
-        "AUTOMATIC1111_API_AUTH": app.state.config.AUTOMATIC1111_API_AUTH,
-        "COMFYUI_BASE_URL": app.state.config.COMFYUI_BASE_URL,
+        "enabled": app.state.config.ENABLED,
+        "engine": app.state.config.ENGINE,
+        "openai": {
+            "OPENAI_API_BASE_URL": app.state.config.OPENAI_API_BASE_URL,
+            "OPENAI_API_KEY": app.state.config.OPENAI_API_KEY,
+        },
+        "automatic1111": {
+            "AUTOMATIC1111_BASE_URL": app.state.config.AUTOMATIC1111_BASE_URL,
+            "AUTOMATIC1111_API_AUTH": app.state.config.AUTOMATIC1111_API_AUTH,
+        },
+        "comfyui": {
+            "COMFYUI_BASE_URL": app.state.config.COMFYUI_BASE_URL,
+            "COMFYUI_WORKFLOW": app.state.config.COMFYUI_WORKFLOW,
+            "COMFYUI_WORKFLOW_NODES": app.state.config.COMFYUI_WORKFLOW_NODES,
+        },
     }
     }
 
 
 
 
-@app.post("/url/update")
-async def update_engine_url(
-    form_data: EngineUrlUpdateForm, user=Depends(get_admin_user)
-):
-    if form_data.AUTOMATIC1111_BASE_URL is None:
-        app.state.config.AUTOMATIC1111_BASE_URL = AUTOMATIC1111_BASE_URL
-    else:
-        url = form_data.AUTOMATIC1111_BASE_URL.strip("/")
-        try:
-            r = requests.head(url)
-            r.raise_for_status()
-            app.state.config.AUTOMATIC1111_BASE_URL = url
-        except Exception as e:
-            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL)
-
-    if form_data.COMFYUI_BASE_URL is None:
-        app.state.config.COMFYUI_BASE_URL = COMFYUI_BASE_URL
-    else:
-        url = form_data.COMFYUI_BASE_URL.strip("/")
-
-        try:
-            r = requests.head(url)
-            r.raise_for_status()
-            app.state.config.COMFYUI_BASE_URL = url
-        except Exception as e:
-            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL)
-
-    if form_data.AUTOMATIC1111_API_AUTH is None:
-        app.state.config.AUTOMATIC1111_API_AUTH = AUTOMATIC1111_API_AUTH
+def get_automatic1111_api_auth():
+    if app.state.config.AUTOMATIC1111_API_AUTH is None:
+        return ""
     else:
     else:
-        app.state.config.AUTOMATIC1111_API_AUTH = form_data.AUTOMATIC1111_API_AUTH
-
-    return {
-        "AUTOMATIC1111_BASE_URL": app.state.config.AUTOMATIC1111_BASE_URL,
-        "AUTOMATIC1111_API_AUTH": app.state.config.AUTOMATIC1111_API_AUTH,
-        "COMFYUI_BASE_URL": app.state.config.COMFYUI_BASE_URL,
-        "status": True,
-    }
+        auth1111_byte_string = app.state.config.AUTOMATIC1111_API_AUTH.encode("utf-8")
+        auth1111_base64_encoded_bytes = base64.b64encode(auth1111_byte_string)
+        auth1111_base64_encoded_string = auth1111_base64_encoded_bytes.decode("utf-8")
+        return f"Basic {auth1111_base64_encoded_string}"
 
 
 
 
-class OpenAIConfigUpdateForm(BaseModel):
-    url: str
-    key: str
+def set_image_model(model: str):
+    app.state.config.MODEL = model
+    if app.state.config.ENGINE in ["", "automatic1111"]:
+        api_auth = get_automatic1111_api_auth()
+        r = requests.get(
+            url=f"{app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
+            headers={"authorization": api_auth},
+        )
+        options = r.json()
+        if model != options["sd_model_checkpoint"]:
+            options["sd_model_checkpoint"] = model
+            r = requests.post(
+                url=f"{app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
+                json=options,
+                headers={"authorization": api_auth},
+            )
+    return app.state.config.MODEL
 
 
 
 
-@app.get("/openai/config")
-async def get_openai_config(user=Depends(get_admin_user)):
-    return {
-        "OPENAI_API_BASE_URL": app.state.config.OPENAI_API_BASE_URL,
-        "OPENAI_API_KEY": app.state.config.OPENAI_API_KEY,
-    }
+def get_image_model():
+    if app.state.config.ENGINE == "openai":
+        return app.state.config.MODEL if app.state.config.MODEL else "dall-e-2"
+    elif app.state.config.ENGINE == "comfyui":
+        return app.state.config.MODEL if app.state.config.MODEL else ""
+    elif app.state.config.ENGINE == "automatic1111" or app.state.config.ENGINE == "":
+        try:
+            r = requests.get(
+                url=f"{app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
+                headers={"authorization": get_automatic1111_api_auth()},
+            )
+            options = r.json()
+            return options["sd_model_checkpoint"]
+        except Exception as e:
+            app.state.config.ENABLED = False
+            raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
 
 
 
 
-@app.post("/openai/config/update")
-async def update_openai_config(
-    form_data: OpenAIConfigUpdateForm, user=Depends(get_admin_user)
-):
-    if form_data.key == "":
-        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)
+class ImageConfigForm(BaseModel):
+    model: str
+    size: str
+    steps: int
 
 
-    app.state.config.OPENAI_API_BASE_URL = form_data.url
-    app.state.config.OPENAI_API_KEY = form_data.key
 
 
+@app.get("/image/config")
+async def get_image_config(user=Depends(get_admin_user)):
     return {
     return {
-        "status": True,
-        "OPENAI_API_BASE_URL": app.state.config.OPENAI_API_BASE_URL,
-        "OPENAI_API_KEY": app.state.config.OPENAI_API_KEY,
+        "MODEL": app.state.config.MODEL,
+        "IMAGE_SIZE": app.state.config.IMAGE_SIZE,
+        "IMAGE_STEPS": app.state.config.IMAGE_STEPS,
     }
     }
 
 
 
 
-class ImageSizeUpdateForm(BaseModel):
-    size: str
-
-
-@app.get("/size")
-async def get_image_size(user=Depends(get_admin_user)):
-    return {"IMAGE_SIZE": app.state.config.IMAGE_SIZE}
-
+@app.post("/image/config/update")
+async def update_image_config(form_data: ImageConfigForm, user=Depends(get_admin_user)):
+    app.state.config.MODEL = form_data.model
 
 
-@app.post("/size/update")
-async def update_image_size(
-    form_data: ImageSizeUpdateForm, user=Depends(get_admin_user)
-):
-    pattern = r"^\d+x\d+$"  # Regular expression pattern
+    pattern = r"^\d+x\d+$"
     if re.match(pattern, form_data.size):
     if re.match(pattern, form_data.size):
         app.state.config.IMAGE_SIZE = form_data.size
         app.state.config.IMAGE_SIZE = form_data.size
-        return {
-            "IMAGE_SIZE": app.state.config.IMAGE_SIZE,
-            "status": True,
-        }
     else:
     else:
         raise HTTPException(
         raise HTTPException(
             status_code=400,
             status_code=400,
             detail=ERROR_MESSAGES.INCORRECT_FORMAT("  (e.g., 512x512)."),
             detail=ERROR_MESSAGES.INCORRECT_FORMAT("  (e.g., 512x512)."),
         )
         )
 
 
-
-class ImageStepsUpdateForm(BaseModel):
-    steps: int
-
-
-@app.get("/steps")
-async def get_image_size(user=Depends(get_admin_user)):
-    return {"IMAGE_STEPS": app.state.config.IMAGE_STEPS}
-
-
-@app.post("/steps/update")
-async def update_image_size(
-    form_data: ImageStepsUpdateForm, user=Depends(get_admin_user)
-):
     if form_data.steps >= 0:
     if form_data.steps >= 0:
         app.state.config.IMAGE_STEPS = form_data.steps
         app.state.config.IMAGE_STEPS = form_data.steps
-        return {
-            "IMAGE_STEPS": app.state.config.IMAGE_STEPS,
-            "status": True,
-        }
     else:
     else:
         raise HTTPException(
         raise HTTPException(
             status_code=400,
             status_code=400,
             detail=ERROR_MESSAGES.INCORRECT_FORMAT("  (e.g., 50)."),
             detail=ERROR_MESSAGES.INCORRECT_FORMAT("  (e.g., 50)."),
         )
         )
 
 
+    return {
+        "MODEL": app.state.config.MODEL,
+        "IMAGE_SIZE": app.state.config.IMAGE_SIZE,
+        "IMAGE_STEPS": app.state.config.IMAGE_STEPS,
+    }
+
 
 
 @app.get("/models")
 @app.get("/models")
 def get_models(user=Depends(get_verified_user)):
 def get_models(user=Depends(get_verified_user)):
@@ -256,7 +262,7 @@ def get_models(user=Depends(get_verified_user)):
                 {"id": "dall-e-3", "name": "DALL·E 3"},
                 {"id": "dall-e-3", "name": "DALL·E 3"},
             ]
             ]
         elif app.state.config.ENGINE == "comfyui":
         elif app.state.config.ENGINE == "comfyui":
-
+            # TODO - get models from comfyui
             r = requests.get(url=f"{app.state.config.COMFYUI_BASE_URL}/object_info")
             r = requests.get(url=f"{app.state.config.COMFYUI_BASE_URL}/object_info")
             info = r.json()
             info = r.json()
 
 
@@ -266,8 +272,9 @@ def get_models(user=Depends(get_verified_user)):
                     info["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0],
                     info["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0],
                 )
                 )
             )
             )
-
-        else:
+        elif (
+            app.state.config.ENGINE == "automatic1111" or app.state.config.ENGINE == ""
+        ):
             r = requests.get(
             r = requests.get(
                 url=f"{app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/sd-models",
                 url=f"{app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/sd-models",
                 headers={"authorization": get_automatic1111_api_auth()},
                 headers={"authorization": get_automatic1111_api_auth()},
@@ -284,69 +291,11 @@ def get_models(user=Depends(get_verified_user)):
         raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
         raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
 
 
 
 
-@app.get("/models/default")
-async def get_default_model(user=Depends(get_admin_user)):
-    try:
-        if app.state.config.ENGINE == "openai":
-            return {
-                "model": (
-                    app.state.config.MODEL if app.state.config.MODEL else "dall-e-2"
-                )
-            }
-        elif app.state.config.ENGINE == "comfyui":
-            return {"model": (app.state.config.MODEL if app.state.config.MODEL else "")}
-        else:
-            r = requests.get(
-                url=f"{app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
-                headers={"authorization": get_automatic1111_api_auth()},
-            )
-            options = r.json()
-            return {"model": options["sd_model_checkpoint"]}
-    except Exception as e:
-        app.state.config.ENABLED = False
-        raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
-
-
-class UpdateModelForm(BaseModel):
-    model: str
-
-
-def set_model_handler(model: str):
-    if app.state.config.ENGINE in ["openai", "comfyui"]:
-        app.state.config.MODEL = model
-        return app.state.config.MODEL
-    else:
-        api_auth = get_automatic1111_api_auth()
-        r = requests.get(
-            url=f"{app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
-            headers={"authorization": api_auth},
-        )
-        options = r.json()
-
-        if model != options["sd_model_checkpoint"]:
-            options["sd_model_checkpoint"] = model
-            r = requests.post(
-                url=f"{app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
-                json=options,
-                headers={"authorization": api_auth},
-            )
-
-        return options
-
-
-@app.post("/models/default/update")
-def update_default_model(
-    form_data: UpdateModelForm,
-    user=Depends(get_verified_user),
-):
-    return set_model_handler(form_data.model)
-
-
 class GenerateImageForm(BaseModel):
 class GenerateImageForm(BaseModel):
     model: Optional[str] = None
     model: Optional[str] = None
     prompt: str
     prompt: str
-    n: int = 1
     size: Optional[str] = None
     size: Optional[str] = None
+    n: int = 1
     negative_prompt: Optional[str] = None
     negative_prompt: Optional[str] = None
 
 
 
 
@@ -497,9 +446,11 @@ async def image_generations(
 
 
             log.debug(f"images: {images}")
             log.debug(f"images: {images}")
             return images
             return images
-        else:
+        elif (
+            app.state.config.ENGINE == "automatic1111" or app.state.config.ENGINE == ""
+        ):
             if form_data.model:
             if form_data.model:
-                set_model_handler(form_data.model)
+                set_image_model(form_data.model)
 
 
             data = {
             data = {
                 "prompt": form_data.prompt,
                 "prompt": form_data.prompt,
@@ -521,7 +472,6 @@ async def image_generations(
             )
             )
 
 
             res = r.json()
             res = r.json()
-
             log.debug(f"res: {res}")
             log.debug(f"res: {res}")
 
 
             images = []
             images = []
@@ -538,7 +488,6 @@ async def image_generations(
 
 
     except Exception as e:
     except Exception as e:
         error = e
         error = e
-
         if r != None:
         if r != None:
             data = r.json()
             data = r.json()
             if "error" in data:
             if "error" in data:

+ 0 - 110
backend/apps/images/utils/comfyui.py

@@ -15,116 +15,6 @@ from pydantic import BaseModel
 
 
 from typing import Optional
 from typing import Optional
 
 
-COMFYUI_DEFAULT_WORKFLOW = """
-{
-  "3": {
-    "inputs": {
-      "seed": 0,
-      "steps": 20,
-      "cfg": 8,
-      "sampler_name": "euler",
-      "scheduler": "normal",
-      "denoise": 1,
-      "model": [
-        "4",
-        0
-      ],
-      "positive": [
-        "6",
-        0
-      ],
-      "negative": [
-        "7",
-        0
-      ],
-      "latent_image": [
-        "5",
-        0
-      ]
-    },
-    "class_type": "KSampler",
-    "_meta": {
-      "title": "KSampler"
-    }
-  },
-  "4": {
-    "inputs": {
-      "ckpt_name": "model.safetensors"
-    },
-    "class_type": "CheckpointLoaderSimple",
-    "_meta": {
-      "title": "Load Checkpoint"
-    }
-  },
-  "5": {
-    "inputs": {
-      "width": 512,
-      "height": 512,
-      "batch_size": 1
-    },
-    "class_type": "EmptyLatentImage",
-    "_meta": {
-      "title": "Empty Latent Image"
-    }
-  },
-  "6": {
-    "inputs": {
-      "text": "Prompt",
-      "clip": [
-        "4",
-        1
-      ]
-    },
-    "class_type": "CLIPTextEncode",
-    "_meta": {
-      "title": "CLIP Text Encode (Prompt)"
-    }
-  },
-  "7": {
-    "inputs": {
-      "text": "",
-      "clip": [
-        "4",
-        1
-      ]
-    },
-    "class_type": "CLIPTextEncode",
-    "_meta": {
-      "title": "CLIP Text Encode (Prompt)"
-    }
-  },
-  "8": {
-    "inputs": {
-      "samples": [
-        "3",
-        0
-      ],
-      "vae": [
-        "4",
-        2
-      ]
-    },
-    "class_type": "VAEDecode",
-    "_meta": {
-      "title": "VAE Decode"
-    }
-  },
-  "9": {
-    "inputs": {
-      "filename_prefix": "ComfyUI",
-      "images": [
-        "8",
-        0
-      ]
-    },
-    "class_type": "SaveImage",
-    "_meta": {
-      "title": "Save Image"
-    }
-  }
-}
-"""
-
 
 
 def queue_prompt(prompt, client_id, base_url):
 def queue_prompt(prompt, client_id, base_url):
     log.info("queue_prompt")
     log.info("queue_prompt")

+ 118 - 1
backend/config.py

@@ -1342,10 +1342,127 @@ COMFYUI_BASE_URL = PersistentConfig(
     os.getenv("COMFYUI_BASE_URL", ""),
     os.getenv("COMFYUI_BASE_URL", ""),
 )
 )
 
 
+COMFYUI_DEFAULT_WORKFLOW = """
+{
+  "3": {
+    "inputs": {
+      "seed": 0,
+      "steps": 20,
+      "cfg": 8,
+      "sampler_name": "euler",
+      "scheduler": "normal",
+      "denoise": 1,
+      "model": [
+        "4",
+        0
+      ],
+      "positive": [
+        "6",
+        0
+      ],
+      "negative": [
+        "7",
+        0
+      ],
+      "latent_image": [
+        "5",
+        0
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "KSampler"
+    }
+  },
+  "4": {
+    "inputs": {
+      "ckpt_name": "model.safetensors"
+    },
+    "class_type": "CheckpointLoaderSimple",
+    "_meta": {
+      "title": "Load Checkpoint"
+    }
+  },
+  "5": {
+    "inputs": {
+      "width": 512,
+      "height": 512,
+      "batch_size": 1
+    },
+    "class_type": "EmptyLatentImage",
+    "_meta": {
+      "title": "Empty Latent Image"
+    }
+  },
+  "6": {
+    "inputs": {
+      "text": "Prompt",
+      "clip": [
+        "4",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "7": {
+    "inputs": {
+      "text": "",
+      "clip": [
+        "4",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "8": {
+    "inputs": {
+      "samples": [
+        "3",
+        0
+      ],
+      "vae": [
+        "4",
+        2
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    }
+  },
+  "9": {
+    "inputs": {
+      "filename_prefix": "ComfyUI",
+      "images": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "SaveImage",
+    "_meta": {
+      "title": "Save Image"
+    }
+  }
+}
+"""
+
+
 COMFYUI_WORKFLOW = PersistentConfig(
 COMFYUI_WORKFLOW = PersistentConfig(
     "COMFYUI_WORKFLOW",
     "COMFYUI_WORKFLOW",
     "image_generation.comfyui.workflow",
     "image_generation.comfyui.workflow",
-    os.getenv("COMFYUI_WORKFLOW", ""),
+    os.getenv("COMFYUI_WORKFLOW", COMFYUI_DEFAULT_WORKFLOW),
+)
+
+COMFYUI_WORKFLOW_NODES = PersistentConfig(
+    "COMFYUI_WORKFLOW",
+    "image_generation.comfyui.nodes",
+    [],
 )
 )
 
 
 IMAGES_OPENAI_API_BASE_URL = PersistentConfig(
 IMAGES_OPENAI_API_BASE_URL = PersistentConfig(

+ 38 - 40
src/lib/components/admin/Settings/Images.svelte

@@ -20,13 +20,14 @@
 	} from '$lib/apis/images';
 	} from '$lib/apis/images';
 	import { getBackendConfig } from '$lib/apis';
 	import { getBackendConfig } from '$lib/apis';
 	import SensitiveInput from '$lib/components/common/SensitiveInput.svelte';
 	import SensitiveInput from '$lib/components/common/SensitiveInput.svelte';
+	import Switch from '$lib/components/common/Switch.svelte';
 	const dispatch = createEventDispatcher();
 	const dispatch = createEventDispatcher();
 
 
 	const i18n = getContext('i18n');
 	const i18n = getContext('i18n');
 
 
 	let loading = false;
 	let loading = false;
 
 
-	let imageGenerationEngine = '';
+	let imageGenerationEngine = 'openai';
 	let enableImageGeneration = false;
 	let enableImageGeneration = false;
 
 
 	let AUTOMATIC1111_BASE_URL = '';
 	let AUTOMATIC1111_BASE_URL = '';
@@ -128,7 +129,7 @@
 			});
 			});
 
 
 			if (res) {
 			if (res) {
-				imageGenerationEngine = res.engine;
+				imageGenerationEngine = res.engine ?? 'automatic1111';
 				enableImageGeneration = res.enabled;
 				enableImageGeneration = res.enabled;
 			}
 			}
 			const URLS = await getImageGenerationEngineUrls(localStorage.token);
 			const URLS = await getImageGenerationEngineUrls(localStorage.token);
@@ -180,6 +181,38 @@
 		<div>
 		<div>
 			<div class=" mb-1 text-sm font-medium">{$i18n.t('Image Settings')}</div>
 			<div class=" mb-1 text-sm font-medium">{$i18n.t('Image Settings')}</div>
 
 
+			<div>
+				<div class=" py-0.5 flex w-full justify-between">
+					<div class=" self-center text-xs font-medium">
+						{$i18n.t('Image Generation (Experimental)')}
+					</div>
+
+					<div class="px-1">
+						<Switch
+							bind:state={enableImageGeneration}
+							on:change={(e) => {
+								const enabled = e.detail;
+
+								if (enabled) {
+									if (imageGenerationEngine === 'automatic1111' && AUTOMATIC1111_BASE_URL === '') {
+										toast.error($i18n.t('AUTOMATIC1111 Base URL is required.'));
+										enableImageGeneration = false;
+									} else if (imageGenerationEngine === 'comfyui' && COMFYUI_BASE_URL === '') {
+										toast.error($i18n.t('ComfyUI Base URL is required.'));
+										enableImageGeneration = false;
+									} else if (imageGenerationEngine === 'openai' && OPENAI_API_KEY === '') {
+										toast.error($i18n.t('OpenAI API Key is required.'));
+										enableImageGeneration = false;
+									}
+								}
+
+								updateImageGeneration();
+							}}
+						/>
+					</div>
+				</div>
+			</div>
+
 			<div class=" py-0.5 flex w-full justify-between">
 			<div class=" py-0.5 flex w-full justify-between">
 				<div class=" self-center text-xs font-medium">{$i18n.t('Image Generation Engine')}</div>
 				<div class=" self-center text-xs font-medium">{$i18n.t('Image Generation Engine')}</div>
 				<div class="flex items-center relative">
 				<div class="flex items-center relative">
@@ -191,51 +224,16 @@
 							await updateImageGeneration();
 							await updateImageGeneration();
 						}}
 						}}
 					>
 					>
-						<option value="">{$i18n.t('Default (Automatic1111)')}</option>
+						<option value="openai">{$i18n.t('Default (Open AI)')}</option>
 						<option value="comfyui">{$i18n.t('ComfyUI')}</option>
 						<option value="comfyui">{$i18n.t('ComfyUI')}</option>
-						<option value="openai">{$i18n.t('Open AI (Dall-E)')}</option>
+						<option value="automatic1111">{$i18n.t('Automatic1111')}</option>
 					</select>
 					</select>
 				</div>
 				</div>
 			</div>
 			</div>
-
-			<div>
-				<div class=" py-0.5 flex w-full justify-between">
-					<div class=" self-center text-xs font-medium">
-						{$i18n.t('Image Generation (Experimental)')}
-					</div>
-
-					<button
-						class="p-1 px-3 text-xs flex rounded transition"
-						on:click={() => {
-							if (imageGenerationEngine === '' && AUTOMATIC1111_BASE_URL === '') {
-								toast.error($i18n.t('AUTOMATIC1111 Base URL is required.'));
-								enableImageGeneration = false;
-							} else if (imageGenerationEngine === 'comfyui' && COMFYUI_BASE_URL === '') {
-								toast.error($i18n.t('ComfyUI Base URL is required.'));
-								enableImageGeneration = false;
-							} else if (imageGenerationEngine === 'openai' && OPENAI_API_KEY === '') {
-								toast.error($i18n.t('OpenAI API Key is required.'));
-								enableImageGeneration = false;
-							} else {
-								enableImageGeneration = !enableImageGeneration;
-							}
-
-							updateImageGeneration();
-						}}
-						type="button"
-					>
-						{#if enableImageGeneration === true}
-							<span class="ml-2 self-center">{$i18n.t('On')}</span>
-						{:else}
-							<span class="ml-2 self-center">{$i18n.t('Off')}</span>
-						{/if}
-					</button>
-				</div>
-			</div>
 		</div>
 		</div>
 		<hr class=" dark:border-gray-850" />
 		<hr class=" dark:border-gray-850" />
 
 
-		{#if imageGenerationEngine === ''}
+		{#if imageGenerationEngine === 'automatic1111'}
 			<div class=" mb-2.5 text-sm font-medium">{$i18n.t('AUTOMATIC1111 Base URL')}</div>
 			<div class=" mb-2.5 text-sm font-medium">{$i18n.t('AUTOMATIC1111 Base URL')}</div>
 			<div class="flex w-full">
 			<div class="flex w-full">
 				<div class="flex-1 mr-2">
 				<div class="flex-1 mr-2">