feat: integrate ComfyUI workflow for local image generation

2025-12-17 21:04:32 +05:45 · 2025-12-17 21:04:32 +05:45 · e72cea3655
commit e72cea3655
parent 0a70f3c4e3
12 changed files with 458 additions and 120 deletions
--- a/docker-compose.yml
+++ b/docker-compose.yml
@ -101,7 +101,7 @@ services:
      - DATABASE_URL=${DATABASE_URL}
      - DISABLE_ANONYMOUS_TRACKING=${DISABLE_ANONYMOUS_TRACKING}
      - LOCAL_IMAGE_URL=${LOCAL_IMAGE_URL}
-      - LOCAL_IMAGE_MODEL=${LOCAL_IMAGE_MODEL}
+      - LOCAL_IMAGE_WORKFLOW=${LOCAL_IMAGE_WORKFLOW}

  development-gpu:
    build:
--- a/servers/fastapi/models/user_config.py
+++ b/servers/fastapi/models/user_config.py
@ -32,9 +32,9 @@ class UserConfig(BaseModel):
    PEXELS_API_KEY: Optional[str] = None
    PIXABAY_API_KEY: Optional[str] = None

-    # Local Image Generation (Stable Diffusion, FLUX, ComfyUI, Fooocus, etc.)
+    # Local Image Generation (ComfyUI)
    LOCAL_IMAGE_URL: Optional[str] = None
-    LOCAL_IMAGE_MODEL: Optional[str] = None
+    LOCAL_IMAGE_WORKFLOW: Optional[str] = None  # ComfyUI workflow JSON

    # Reasoning
    TOOL_CALLS: Optional[bool] = None
--- a/servers/fastapi/services/image_generation_service.py
+++ b/servers/fastapi/services/image_generation_service.py
@ -1,5 +1,6 @@
 import asyncio
 import base64
+import json
 import os
 import aiohttp
 from google import genai
@ -11,7 +12,7 @@ from utils.download_helpers import download_file
 from utils.get_env import get_pexels_api_key_env
 from utils.get_env import get_pixabay_api_key_env
 from utils.get_env import get_local_image_url_env
-from utils.get_env import get_local_image_model_env
+from utils.get_env import get_local_image_workflow_env
 from utils.image_provider import (
    is_image_generation_disabled,
    is_pixels_selected,
@ -146,17 +147,14 @@ class ImageGenerationService:

    async def generate_image_local(self, prompt: str, output_directory: str) -> str:
        """
-        Generate image using a local image generation server.
+        Generate image using ComfyUI workflow API.
        
-        User provides the full API URL including the endpoint.
-        Examples:
-        - Automatic1111: http://192.168.1.7:7860/sdapi/v1/txt2img
-        - Fooocus: http://192.168.1.7:7860/v1/generation/text-to-image
-        - Custom: http://192.168.1.7:7860/generate
+        User provides:
+        - LOCAL_IMAGE_URL: ComfyUI server URL (e.g., http://192.168.1.7:8188)
+        - LOCAL_IMAGE_WORKFLOW: Workflow JSON exported from ComfyUI
        
-        Supports both:
-        - JSON response with base64 images (Automatic1111 style)
-        - Direct binary image response (raw PNG/JPEG)
+        The workflow should have a CLIPTextEncode node with "Positive" in the title
+        where the prompt will be injected.
        
        Args:
            prompt: The text prompt for image generation
@ -165,82 +163,205 @@ class ImageGenerationService:
        Returns:
            Path to the generated image file
        """
-        api_url = get_local_image_url_env()
-        local_model = get_local_image_model_env()
+        comfyui_url = get_local_image_url_env()
+        workflow_json = get_local_image_workflow_env()
        
-        if not api_url:
+        if not comfyui_url:
            raise ValueError("LOCAL_IMAGE_URL environment variable is not set")
        
-        # Build the request payload (Automatic1111 compatible format)
-        # Most local tools accept similar payload structure
-        payload = {
-            "prompt": prompt,
-            "negative_prompt": "blurry, bad quality, distorted, ugly, deformed",
-            "steps": 20,
-            "width": 1024,
-            "height": 1024,
-            "cfg_scale": 7,
-            "sampler_name": "Euler a",
-        }
+        if not workflow_json:
+            raise ValueError("LOCAL_IMAGE_WORKFLOW environment variable is not set. Please provide a ComfyUI workflow JSON.")
        
-        # Add model override if specified
-        if local_model:
-            payload["override_settings"] = {
-                "sd_model_checkpoint": local_model
-            }
+        # Ensure URL doesn't have trailing slash
+        comfyui_url = comfyui_url.rstrip("/")
+        
+        # Parse the workflow JSON
+        try:
+            workflow = json.loads(workflow_json)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid workflow JSON: {str(e)}")
+        
+        # Find and update the positive prompt node
+        workflow = self._inject_prompt_into_workflow(workflow, prompt)
        
        async with aiohttp.ClientSession(trust_env=True) as session:
+            # Step 1: Submit workflow
+            prompt_id = await self._submit_comfyui_workflow(session, comfyui_url, workflow)
+            
+            # Step 2: Wait for completion
+            status_data = await self._wait_for_comfyui_completion(session, comfyui_url, prompt_id)
+            
+            # Step 3: Download the generated image
+            image_path = await self._download_comfyui_image(
+                session, comfyui_url, status_data, prompt_id, output_directory
+            )
+            
+            return image_path
+    
+    def _inject_prompt_into_workflow(self, workflow: dict, prompt: str) -> dict:
+        """
+        Find the positive prompt node in the workflow and inject the prompt text.
+        Looks for CLIPTextEncode nodes with 'Positive' in the title.
+        """
+        prompt_injected = False
+        
+        for node_id, node_data in workflow.items():
+            # Check if this is a CLIPTextEncode node
+            if node_data.get("class_type") == "CLIPTextEncode":
+                meta = node_data.get("_meta", {})
+                title = meta.get("title", "").lower()
+                
+                # Check if it's a positive prompt node
+                if "positive" in title:
+                    if "inputs" in node_data and "text" in node_data["inputs"]:
+                        node_data["inputs"]["text"] = prompt
+                        prompt_injected = True
+                        print(f"Injected prompt into node {node_id}: {title}")
+                        break
+        
+        if not prompt_injected:
+            # Fallback: try to find any CLIPTextEncode node with text input
+            for node_id, node_data in workflow.items():
+                if node_data.get("class_type") == "CLIPTextEncode":
+                    if "inputs" in node_data and "text" in node_data["inputs"]:
+                        # Skip if it looks like a negative prompt
+                        meta = node_data.get("_meta", {})
+                        title = meta.get("title", "").lower()
+                        if "negative" in title:
+                            continue
+                        node_data["inputs"]["text"] = prompt
+                        prompt_injected = True
+                        print(f"Injected prompt into node {node_id} (fallback)")
+                        break
+        
+        if not prompt_injected:
+            raise ValueError("Could not find a positive prompt node (CLIPTextEncode) in the workflow")
+        
+        return workflow
+    
+    async def _submit_comfyui_workflow(
+        self, session: aiohttp.ClientSession, comfyui_url: str, workflow: dict
+    ) -> str:
+        """Submit workflow to ComfyUI and return the prompt_id."""
+        client_id = str(uuid.uuid4())
+        payload = {
+            "prompt": workflow,
+            "client_id": client_id
+        }
+        
+        response = await session.post(
+            f"{comfyui_url}/prompt",
+            json=payload,
+            timeout=aiohttp.ClientTimeout(total=30)
+        )
+        
+        if response.status != 200:
+            error_text = await response.text()
+            raise Exception(f"Failed to submit workflow to ComfyUI: {error_text}")
+        
+        data = await response.json()
+        prompt_id = data.get("prompt_id")
+        
+        if not prompt_id:
+            raise Exception("No prompt_id returned from ComfyUI")
+        
+        print(f"ComfyUI workflow submitted. Prompt ID: {prompt_id}")
+        return prompt_id
+    
+    async def _wait_for_comfyui_completion(
+        self, session: aiohttp.ClientSession, comfyui_url: str, prompt_id: str,
+        timeout: int = 300, poll_interval: int = 4
+    ) -> dict:
+        """Poll ComfyUI history endpoint until workflow completes."""
+        start_time = asyncio.get_event_loop().time()
+        
+        while True:
+            elapsed = asyncio.get_event_loop().time() - start_time
+            if elapsed > timeout:
+                raise Exception(f"ComfyUI workflow timed out after {timeout} seconds")
+            
+            await asyncio.sleep(poll_interval)
+            
+            response = await session.get(
+                f"{comfyui_url}/history/{prompt_id}",
+                timeout=aiohttp.ClientTimeout(total=30)
+            )
+            
+            if response.status != 200:
+                continue
+            
            try:
-                response = await session.post(
-                    api_url,
-                    json=payload,
-                    timeout=aiohttp.ClientTimeout(total=300)  # 5 min timeout for generation
-                )
+                status_data = await response.json()
+            except:
+                continue
+            
+            if prompt_id in status_data:
+                execution_data = status_data[prompt_id]
                
-                if response.status != 200:
-                    error_text = await response.text()
-                    raise Exception(f"Local image API error: {response.status} - {error_text}")
+                # Check for completion
+                if "status" in execution_data:
+                    status = execution_data["status"]
+                    if status.get("completed", False):
+                        print("ComfyUI workflow completed successfully")
+                        return status_data
+                    if "error" in status:
+                        raise Exception(f"ComfyUI workflow error: {status['error']}")
                
-                content_type = response.headers.get("Content-Type", "")
-                
-                # Handle direct binary image response (image/png, image/jpeg, etc.)
-                if content_type.startswith("image/"):
-                    image_data = await response.read()
-                    # Determine file extension from content type
-                    ext = "png" if "png" in content_type else "jpg"
-                    image_path = os.path.join(output_directory, f"{uuid.uuid4()}.{ext}")
+                # Also check if outputs exist (alternative completion check)
+                if "outputs" in execution_data and execution_data["outputs"]:
+                    print("ComfyUI workflow completed (outputs found)")
+                    return status_data
+            
+            print(f"Waiting for ComfyUI workflow... ({int(elapsed)}s)")
+    
+    async def _download_comfyui_image(
+        self, session: aiohttp.ClientSession, comfyui_url: str,
+        status_data: dict, prompt_id: str, output_directory: str
+    ) -> str:
+        """Download the generated image from ComfyUI."""
+        if prompt_id not in status_data:
+            raise Exception("Prompt ID not found in status data")
+        
+        outputs = status_data[prompt_id].get("outputs", {})
+        
+        if not outputs:
+            raise Exception("No outputs found in ComfyUI response")
+        
+        # Find the first image in outputs
+        for node_id, node_output in outputs.items():
+            if "images" in node_output:
+                for image_info in node_output["images"]:
+                    filename = image_info["filename"]
+                    subfolder = image_info.get("subfolder", "")
                    
-                    with open(image_path, "wb") as f:
-                        f.write(image_data)
+                    # Build view params
+                    params = {
+                        "filename": filename,
+                        "type": "output"
+                    }
+                    if subfolder:
+                        params["subfolder"] = subfolder
                    
-                    return image_path
-                
-                # Handle JSON response with base64 encoded images
-                data = await response.json()
-                
-                # Check for images in various response formats
-                if "images" in data and len(data["images"]) > 0:
-                    image_base64 = data["images"][0]
-                    # Handle if it's a dict with base64 key
-                    if isinstance(image_base64, dict) and "base64" in image_base64:
-                        image_base64 = image_base64["base64"]
-                elif "image" in data:
-                    image_base64 = data["image"]
-                elif "output" in data:
-                    image_base64 = data["output"]
-                elif "result" in data:
-                    image_base64 = data["result"]
-                else:
-                    raise Exception(f"No images found in response. Keys: {list(data.keys())}")
-                
-                # Decode base64 and save to file
-                image_data = base64.b64decode(image_base64)
-                image_path = os.path.join(output_directory, f"{uuid.uuid4()}.png")
-                
-                with open(image_path, "wb") as f:
-                    f.write(image_data)
-                
-                return image_path
+                    # Download the image
+                    response = await session.get(
+                        f"{comfyui_url}/view",
+                        params=params,
+                        timeout=aiohttp.ClientTimeout(total=60)
+                    )
                    
-            except aiohttp.ClientError as e:
-                raise Exception(f"Failed to connect to local image server at {api_url}: {str(e)}")
+                    if response.status == 200:
+                        image_data = await response.read()
+                        
+                        # Determine extension
+                        ext = filename.split(".")[-1] if "." in filename else "png"
+                        image_path = os.path.join(output_directory, f"{uuid.uuid4()}.{ext}")
+                        
+                        with open(image_path, "wb") as f:
+                            f.write(image_data)
+                        
+                        print(f"Downloaded image from ComfyUI: {image_path}")
+                        return image_path
+                    else:
+                        raise Exception(f"Failed to download image: {response.status}")
+        
+        raise Exception("No images found in ComfyUI outputs")
--- a/servers/fastapi/utils/get_env.py
+++ b/servers/fastapi/utils/get_env.py
@ -105,5 +105,5 @@ def get_local_image_url_env():
    return os.getenv("LOCAL_IMAGE_URL")


-def get_local_image_model_env():
-    return os.getenv("LOCAL_IMAGE_MODEL")
+def get_local_image_workflow_env():
+    return os.getenv("LOCAL_IMAGE_WORKFLOW")
--- a/servers/fastapi/utils/set_env.py
+++ b/servers/fastapi/utils/set_env.py
@ -93,5 +93,5 @@ def set_local_image_url_env(value):
    os.environ["LOCAL_IMAGE_URL"] = value


-def set_local_image_model_env(value):
-    os.environ["LOCAL_IMAGE_MODEL"] = value
+def set_local_image_workflow_env(value):
+    os.environ["LOCAL_IMAGE_WORKFLOW"] = value
--- a/servers/fastapi/utils/user_config.py
+++ b/servers/fastapi/utils/user_config.py
@ -13,8 +13,8 @@ from utils.get_env import (
    get_google_api_key_env,
    get_google_model_env,
    get_llm_provider_env,
-    get_local_image_model_env,
    get_local_image_url_env,
+    get_local_image_workflow_env,
    get_ollama_model_env,
    get_ollama_url_env,
    get_openai_api_key_env,
@ -40,8 +40,8 @@ from utils.set_env import (
    set_google_api_key_env,
    set_google_model_env,
    set_llm_provider_env,
-    set_local_image_model_env,
    set_local_image_url_env,
+    set_local_image_workflow_env,
    set_ollama_model_env,
    set_ollama_url_env,
    set_openai_api_key_env,
@ -90,7 +90,7 @@ def get_user_config():
        PIXABAY_API_KEY=existing_config.PIXABAY_API_KEY or get_pixabay_api_key_env(),
        PEXELS_API_KEY=existing_config.PEXELS_API_KEY or get_pexels_api_key_env(),
        LOCAL_IMAGE_URL=existing_config.LOCAL_IMAGE_URL or get_local_image_url_env(),
-        LOCAL_IMAGE_MODEL=existing_config.LOCAL_IMAGE_MODEL or get_local_image_model_env(),
+        LOCAL_IMAGE_WORKFLOW=existing_config.LOCAL_IMAGE_WORKFLOW or get_local_image_workflow_env(),
        TOOL_CALLS=(
            existing_config.TOOL_CALLS
            if existing_config.TOOL_CALLS is not None
@ -150,8 +150,8 @@ def update_env_with_user_config():
        set_pexels_api_key_env(user_config.PEXELS_API_KEY)
    if user_config.LOCAL_IMAGE_URL:
        set_local_image_url_env(user_config.LOCAL_IMAGE_URL)
-    if user_config.LOCAL_IMAGE_MODEL:
-        set_local_image_model_env(user_config.LOCAL_IMAGE_MODEL)
+    if user_config.LOCAL_IMAGE_WORKFLOW:
+        set_local_image_workflow_env(user_config.LOCAL_IMAGE_WORKFLOW)
    if user_config.TOOL_CALLS is not None:
        set_tool_calls_env(str(user_config.TOOL_CALLS))
    if user_config.DISABLE_THINKING is not None:
--- a/servers/nextjs/app/api/user-config/route.ts
+++ b/servers/nextjs/app/api/user-config/route.ts
@ -65,7 +65,7 @@ export async function POST(request: Request) {
    IMAGE_PROVIDER: userConfig.IMAGE_PROVIDER || existingConfig.IMAGE_PROVIDER,
    PEXELS_API_KEY: userConfig.PEXELS_API_KEY || existingConfig.PEXELS_API_KEY,
    LOCAL_IMAGE_URL: userConfig.LOCAL_IMAGE_URL || existingConfig.LOCAL_IMAGE_URL,
-    LOCAL_IMAGE_MODEL: userConfig.LOCAL_IMAGE_MODEL || existingConfig.LOCAL_IMAGE_MODEL,
+    LOCAL_IMAGE_WORKFLOW: userConfig.LOCAL_IMAGE_WORKFLOW || existingConfig.LOCAL_IMAGE_WORKFLOW,
    TOOL_CALLS:
      userConfig.TOOL_CALLS === undefined
        ? existingConfig.TOOL_CALLS
--- a/servers/nextjs/components/LLMSelection.tsx
+++ b/servers/nextjs/components/LLMSelection.tsx
@ -82,13 +82,14 @@ export default function LLMProviderSelection({

    const needsOllamaUrl = (llmConfig.LLM === "ollama" && !llmConfig.OLLAMA_URL);

-    const needsLocalImageUrl = !llmConfig.DISABLE_IMAGE_GENERATION &&
-      llmConfig.IMAGE_PROVIDER === "local" && !llmConfig.LOCAL_IMAGE_URL;
+    const needsComfyUIConfig = !llmConfig.DISABLE_IMAGE_GENERATION &&
+      llmConfig.IMAGE_PROVIDER === "local" &&
+      (!llmConfig.LOCAL_IMAGE_URL || !llmConfig.LOCAL_IMAGE_WORKFLOW);

    setButtonState({
      isLoading: false,
-      isDisabled: needsModelSelection || needsApiKey || needsOllamaUrl || needsLocalImageUrl,
-      text: needsModelSelection ? "Please Select a Model" : needsApiKey ? "Please Enter API Key" : needsOllamaUrl ? "Please Enter Ollama URL" : needsLocalImageUrl ? "Please Enter Local Server URL" : "Save Configuration",
+      isDisabled: needsModelSelection || needsApiKey || needsOllamaUrl || needsComfyUIConfig,
+      text: needsModelSelection ? "Please Select a Model" : needsApiKey ? "Please Enter API Key" : needsOllamaUrl ? "Please Enter Ollama URL" : needsComfyUIConfig ? "Please Configure ComfyUI" : "Save Configuration",
      showProgress: false
    });

@ -339,18 +340,18 @@ export default function LLMProviderSelection({
                  return <></>;
                }

-                // Show Local Image Generation configuration
+                // Show ComfyUI configuration
                if (provider.value === "local") {
                  return (
                    <div className="mb-8 space-y-4">
                      <div>
                        <label className="block text-sm font-medium text-gray-700 mb-2">
-                          Local API URL (Full Endpoint)
+                          ComfyUI Server URL
                        </label>
                        <div className="relative">
                          <input
                            type="text"
-                            placeholder="http://192.168.1.7:7860/sdapi/v1/txt2img"
+                            placeholder="http://192.168.1.7:8188"
                            className="w-full px-4 py-2.5 outline-none border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500/20 focus:border-blue-500 transition-colors"
                            value={llmConfig.LOCAL_IMAGE_URL || ""}
                            onChange={(e) => {
@ -358,33 +359,29 @@ export default function LLMProviderSelection({
                            }}
                          />
                        </div>
-                        <p className="mt-2 text-sm text-gray-500">
-                          Enter the full API URL including endpoint. Examples:
+                        <p className="mt-2 text-sm text-gray-500 flex items-center gap-2">
+                          <span className="block w-1 h-1 rounded-full bg-gray-400"></span>
+                          Use your machine IP address (not localhost) when running in Docker
                        </p>
-                        <ul className="mt-1 text-xs text-gray-500 space-y-0.5 ml-4">
-                          <li>• Automatic1111: <code className="bg-gray-100 px-1 rounded">http://IP:7860/sdapi/v1/txt2img</code></li>
-                          <li>• Fooocus: <code className="bg-gray-100 px-1 rounded">http://IP:7860/v1/generation/text-to-image</code></li>
-                          <li>• Use your machine IP address, not localhost</li>
-                        </ul>
                      </div>
                      <div>
                        <label className="block text-sm font-medium text-gray-700 mb-2">
-                          Model Checkpoint (Optional)
+                          Workflow JSON
                        </label>
                        <div className="relative">
-                          <input
-                            type="text"
-                            placeholder="e.g., sd_xl_base_1.0.safetensors or flux1-dev.safetensors"
-                            className="w-full px-4 py-2.5 outline-none border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500/20 focus:border-blue-500 transition-colors"
-                            value={llmConfig.LOCAL_IMAGE_MODEL || ""}
+                          <textarea
+                            placeholder='Paste your ComfyUI workflow JSON here (export via "Save (API Format)" in ComfyUI)'
+                            className="w-full px-4 py-2.5 outline-none border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500/20 focus:border-blue-500 transition-colors font-mono text-xs"
+                            rows={6}
+                            value={llmConfig.LOCAL_IMAGE_WORKFLOW || ""}
                            onChange={(e) => {
-                              input_field_changed(e.target.value, "local_image_model");
+                              input_field_changed(e.target.value, "local_image_workflow");
                            }}
                          />
                        </div>
-                        <p className="mt-2 text-sm text-gray-500 flex items-center gap-2">
-                          <span className="block w-1 h-1 rounded-full bg-gray-400"></span>
-                          Leave empty to use the currently loaded model
+                        <p className="mt-2 text-sm text-gray-500">
+                          Export your workflow from ComfyUI using &quot;Save (API Format)&quot; and paste the JSON here.
+                          The positive prompt node (CLIPTextEncode) will be automatically updated.
                        </p>
                      </div>
                    </div>
--- a/servers/nextjs/types/llm_config.ts
+++ b/servers/nextjs/types/llm_config.ts
@ -28,9 +28,9 @@ export interface LLMConfig {
  PEXELS_API_KEY?: string;
  PIXABAY_API_KEY?: string;

-  // Local Image Generation (Stable Diffusion, FLUX, ComfyUI, Fooocus, etc.)
+  // Local Image Generation (ComfyUI)
  LOCAL_IMAGE_URL?: string;
-  LOCAL_IMAGE_MODEL?: string;
+  LOCAL_IMAGE_WORKFLOW?: string;  // ComfyUI workflow JSON

  // Other Configs
  TOOL_CALLS?: boolean;
--- a/servers/nextjs/utils/providerConstants.ts
+++ b/servers/nextjs/utils/providerConstants.ts
@ -63,12 +63,12 @@ export const IMAGE_PROVIDERS: Record<string, ImageProviderOption> = {
  },
  local: {
    value: "local",
-    label: "Local Image Generation",
-    description: "Use your local AI image server (Stable Diffusion, FLUX, ComfyUI, Fooocus, etc.)",
+    label: "ComfyUI",
+    description: "Use your local ComfyUI server with custom workflows",
    icon: "/icons/local.png",
    requiresApiKey: false,
    apiKeyField: "LOCAL_IMAGE_URL",
-    apiKeyFieldLabel: "Local Server URL"
+    apiKeyFieldLabel: "ComfyUI Server URL"
  },
 };

--- a/servers/nextjs/utils/providerUtils.ts
+++ b/servers/nextjs/utils/providerUtils.ts
@ -49,7 +49,7 @@ export const updateLLMConfig = (
    extended_reasoning: "EXTENDED_REASONING",
    web_grounding: "WEB_GROUNDING",
    local_image_url: "LOCAL_IMAGE_URL",
-    local_image_model: "LOCAL_IMAGE_MODEL",
+    local_image_workflow: "LOCAL_IMAGE_WORKFLOW",
  };

  const configKey = fieldMappings[field];
--- a/text.py
+++ b/text.py
@ -0,0 +1,220 @@
+import json
+import uuid
+import requests
+import time
+
+
+POSITIVE_PROMPT = "Modern abstract representation of global short-form video marketing solutions, showing professional data flows and integrated mobile screens displaying diverse video content, utilizing a neutral and earthy color palette with high contrast. Clean lines and geometric shapes dominate the composition, evoking a sense of technological sophistication and connectivity. The artwork should convey the dynamic and fast-paced nature of short-form video marketing, with an emphasis on innovation and digital communication. The style should be sleek and contemporary, suitable for a corporate audience interested in cutting-edge marketing strategies."
+
+COMFYUI_URL = "https://qfrtn6he9wnwog-8188.proxy.runpod.net"
+
+
+
+
+
+def create_workflow():
+    workflow= {
+  "6": {
+    "inputs": {
+      "text": POSITIVE_PROMPT,
+      "clip": [
+        "30",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Positive Prompt)"
+    }
+  },
+  "8": {
+    "inputs": {
+      "samples": [
+        "31",
+        0
+      ],
+      "vae": [
+        "30",
+        2
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    }
+  },
+  "9": {
+    "inputs": {
+      "filename_prefix": "ComfyUI",
+      "images": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "SaveImage",
+    "_meta": {
+      "title": "Save Image"
+    }
+  },
+  "27": {
+    "inputs": {
+      "width": 1024,
+      "height": 1024,
+      "batch_size": 1
+    },
+    "class_type": "EmptySD3LatentImage",
+    "_meta": {
+      "title": "EmptySD3LatentImage"
+    }
+  },
+  "30": {
+    "inputs": {
+      "ckpt_name": 'flux1-schnell-fp8.safetensors'
+    },
+    "class_type": "CheckpointLoaderSimple",
+    "_meta": {
+      "title": "Load Checkpoint"
+    }
+  },
+  "31": {
+    "inputs": {
+      "seed": 5542493640978,
+      "steps": 4,
+      "cfg": 1,
+      "sampler_name": "euler",
+      "scheduler": "simple",
+      "denoise": 1,
+      "model": [
+        "30",
+        0
+      ],
+      "positive": [
+        "6",
+        0
+      ],
+      "negative": [
+        "33",
+        0
+      ],
+      "latent_image": [
+        "27",
+        0
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "KSampler"
+    }
+  },
+  "33": {
+    "inputs": {
+      "text": "",
+      "clip": [
+        "30",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Negative Prompt)"
+    }
+  }
+}
+    return workflow
+
+def generate_client_id():
+    return str(uuid.uuid4())
+
+def submit_workflow():
+    workflow = create_workflow()
+    client_id = generate_client_id()
+    payload = {
+        "prompt":workflow,
+        "client_id":client_id
+    }
+    print(f"sending request to comfyui...{COMFYUI_URL}/prompt submit_workflow")
+    response = requests.post(f"{COMFYUI_URL}/prompt", json=payload)
+    if response.status_code != 200:
+        raise Exception(f"Failed to submit workflow: {response.text}")
+    
+    response_data = response.json()
+    prompt_id = response_data.get("prompt_id")
+    print(f"Workflow submitted successfully. Prompt ID: {prompt_id}")
+    return prompt_id
+
+def wait_for_completion(prompt_id):
+    print("Waiting for workflow to complete...")
+    while True:
+        time.sleep(5)
+        response = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
+        if response.status_code != 200:
+            raise Exception(f"Failed to get workflow status: {response.text}")
+        try:
+            status_data = response.json()
+           
+        except json.JSONDecodeError:
+            print("Received invalid JSON response, retrying...")
+            continue
+        if prompt_id in status_data:
+            execution_data = status_data[prompt_id] 
+            if 'status' in execution_data and execution_data['status'].get('completed', False):
+                print("Workflow completed.")
+                return status_data
+            if 'status' in execution_data and 'error' in execution_data['status']:
+                print(f"Workflow error: {execution_data['status']['error']}")
+                return None
+        print("Workflow not completed yet, checking again...")
+        
+def get_image_url(status_data, prompt_id):
+    if prompt_id not in status_data or 'outputs' not in status_data[prompt_id]:
+        print("No outputs found for the given prompt ID.")
+        return
+    outputs = status_data[prompt_id]['outputs'] 
+    images_downloaded=0
+    
+    for node_id, node_output in outputs.items():
+        if 'images' in node_output:
+            for image_info in node_output['images']:
+                filename = image_info['filename']
+                subfolder = image_info.get('subfolder', '')
+                
+                view_params ={
+                    "filename": filename,
+                    "type": "output",
+                }
+                if subfolder:
+                    view_params["subfolder"] = subfolder
+                print('downloading image:', filename)
+                image_response = requests.get(f"{COMFYUI_URL}/view", params=view_params)
+                
+                if image_response.status_code == 200:
+                    output_filename = f"output_{images_downloaded}_{filename}"
+                    with open(output_filename, 'wb') as image_file:
+                        image_file.write(image_response.content)
+                    print(f"Image saved as {output_filename}")
+                    
+                    images_downloaded
+                else:
+                    print(f"Failed to download image {filename}: {image_response.text}")
+    if images_downloaded == 0:
+        print("No images were downloaded.")
+    else:
+        print(f"Total images downloaded: {images_downloaded}")
+        
+def main():
+    print("Starting workflow submission...")
+    prompt_id = submit_workflow()
+    
+    if not prompt_id:
+        print("Failed to submit workflow.")
+        return
+    status_data = wait_for_completion(prompt_id=prompt_id)
+    if not status_data:
+        print("Workflow execution failed.")
+        return
+    get_image_url(status_data=status_data, prompt_id=prompt_id)
+    
+    
+    
+if __name__ == "__main__":
+    main()