ComfyUI workflow for 2 images

external_workers/llm_workers/blender_worker.py

@@ -1,13 +1,15 @@
 import asyncio
 import os
 import logging
+import uuid
 from typing import Optional
 import re
 
 import httpx
 from camunda.external_task.external_task import ExternalTask, TaskResult
 from camunda.external_task.external_task_worker import ExternalTaskWorker
 
+from llm_workers.gen_image_workflow import ImageGenerator
 from llm_workers.openai_describe import describe_image_with_openai_vision, name_description_based_of_vision_description
 from season_pass_invite.config import SYS_KEY, API_URL, GEN_IMG_URL, OLLAMA_URL
 
@@ -21,6 +23,10 @@
 X_ACCOUNT = os.getenv('X_ACCOUNT', 'xgurunetwork')
 X_LINK = os.getenv('X_LINK', 'https://v2.dex.guru/gen/')
 
+COMFY_SERVER_HOST = os.getenv('COMFY_SERVER_HOST', 'localhost:8188')
+PROMPT_FILE_PATH = os.getenv('PROMPT_FILE_PATH', 'llm_workers/comfy_workflows/2_images_workflow.json')
+
+
 # Logging the script startup
 logging.info("Starting Image Generation Worker")
 
@@ -34,6 +40,10 @@
     "sleepSeconds": 10
 }
 
+image_gen = ImageGenerator(server_address=COMFY_SERVER_HOST,
+                           prompt_file_path=PROMPT_FILE_PATH)
+
+
 async def fetch_image_url(art_id: str) -> Optional[str]:
     async with httpx.AsyncClient() as _client:
         headers = {
@@ -138,8 +148,14 @@ async def generate_and_upload_image(src1_art_id: str, src2_art_id: str, camunda_
         fetch_image_url(src2_art_id),
         get_user_id_by_camunda_user_id(camunda_user_id)
     )
-    s3_url = await gen_image_client(image1_url, image2_url)
 
+    output_filename = f"gen_img/{uuid.uuid4()}.png"
+
+    result = await image_gen.gen_image(src_image_url_1=image1_url, src_image_url_2=image2_url,
+                                       image_name=output_filename)
+    if result.get("error"):
+        raise Exception("Error while generating images")
+    s3_url = result["image_url"]
 
     visual_description = await describe_image_with_openai_vision(s3_url, "default", "default",
                                                                  "generated_art",)

external_workers/llm_workers/comfy_workflows/2_images_workflow.json

@@ -0,0 +1,230 @@
+{
+  "1": {
+    "inputs": {
+      "ckpt_name": "sd_xl_base_1.0.safetensors"
+    },
+    "class_type": "CheckpointLoaderSimple",
+    "_meta": {
+      "title": "Load Checkpoint"
+    }
+  },
+  "2": {
+    "inputs": {
+      "clip_name": "clip_vision_g.safetensors"
+    },
+    "class_type": "CLIPVisionLoader",
+    "_meta": {
+      "title": "Load CLIP Vision"
+    }
+  },
+  "3": {
+    "inputs": {
+      "text": "Meditative character, vibrant hues, blue, green, orange, swirling patterns, robots, stormtroopers, dynamic poses, harmony, contrast, light, shadow, depth, tranquility, wonder, high-value NFT art",
+      "clip": [
+        "1",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "4": {
+    "inputs": {
+      "text": "",
+      "clip": [
+        "1",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "5": {
+    "inputs": {
+      "conditioning": [
+        "3",
+        0
+      ]
+    },
+    "class_type": "ConditioningZeroOut",
+    "_meta": {
+      "title": "ConditioningZeroOut"
+    }
+  },
+  "6": {
+    "inputs": {
+      "conditioning": [
+        "4",
+        0
+      ]
+    },
+    "class_type": "ConditioningZeroOut",
+    "_meta": {
+      "title": "ConditioningZeroOut"
+    }
+  },
+  "9": {
+    "inputs": {
+      "clip_vision": [
+        "2",
+        0
+      ],
+      "image": [
+        "31",
+        0
+      ]
+    },
+    "class_type": "CLIPVisionEncode",
+    "_meta": {
+      "title": "CLIP Vision Encode"
+    }
+  },
+  "10": {
+    "inputs": {
+      "clip_vision": [
+        "2",
+        0
+      ],
+      "image": [
+        "32",
+        0
+      ]
+    },
+    "class_type": "CLIPVisionEncode",
+    "_meta": {
+      "title": "CLIP Vision Encode"
+    }
+  },
+  "11": {
+    "inputs": {
+      "strength": 1,
+      "noise_augmentation": 0,
+      "conditioning": [
+        "5",
+        0
+      ],
+      "clip_vision_output": [
+        "9",
+        0
+      ]
+    },
+    "class_type": "unCLIPConditioning",
+    "_meta": {
+      "title": "unCLIPConditioning"
+    }
+  },
+  "12": {
+    "inputs": {
+      "strength": 1,
+      "noise_augmentation": 0,
+      "conditioning": [
+        "11",
+        0
+      ],
+      "clip_vision_output": [
+        "10",
+        0
+      ]
+    },
+    "class_type": "unCLIPConditioning",
+    "_meta": {
+      "title": "unCLIPConditioning"
+    }
+  },
+  "13": {
+    "inputs": {
+      "seed": 827367775199718,
+      "steps": 25,
+      "cfg": 8,
+      "sampler_name": "euler",
+      "scheduler": "normal",
+      "denoise": 1,
+      "model": [
+        "1",
+        0
+      ],
+      "positive": [
+        "12",
+        0
+      ],
+      "negative": [
+        "6",
+        0
+      ],
+      "latent_image": [
+        "16",
+        0
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "KSampler"
+    }
+  },
+  "16": {
+    "inputs": {
+      "width": 1024,
+      "height": 1024,
+      "batch_size": 1
+    },
+    "class_type": "EmptyLatentImage",
+    "_meta": {
+      "title": "Empty Latent Image"
+    }
+  },
+  "17": {
+    "inputs": {
+      "samples": [
+        "13",
+        0
+      ],
+      "vae": [
+        "1",
+        2
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    }
+  },
+  "18": {
+    "inputs": {
+      "filename_prefix": "ComfyUI",
+      "images": [
+        "17",
+        0
+      ]
+    },
+    "class_type": "SaveImageWebsocket",
+    "_meta": {
+      "title": "SaveImageWebsocket"
+    }
+  },
+  "31": {
+    "inputs": {
+      "image_url": "https://assets-stage.dex.guru/season_2/maskots/maskots05.png",
+      "save_file_name_override": "",
+      "save_path": ""
+    },
+    "class_type": "🖼️ Download Image from URL",
+    "_meta": {
+      "title": "🖼️ Download Image from URL"
+    }
+  },
+  "32": {
+    "inputs": {
+      "image_url": "https://assets-stage.dex.guru/season_2/seasons/seasonpass_01.jpg",
+      "save_file_name_override": "",
+      "save_path": ""
+    },
+    "class_type": "🖼️ Download Image from URL",
+    "_meta": {
+      "title": "🖼️ Download Image from URL"
+    }
+  }
+}

external_workers/llm_workers/gen_image_workflow.py

@@ -0,0 +1,86 @@
+import asyncio
+import json
+import logging
+import urllib.parse
+import urllib.request
+import uuid
+
+import websocket
+
+from llm_workers.models_server.utils import upload_file_to_s3_binary
+
+
+class ImageGenerator:
+    def __init__(self, server_address,
+                 prompt_file_path):
+        self.server_address = server_address
+        self.client_id = str(uuid.uuid4())
+        self.prompt_data = self.load_prompt(prompt_file_path)
+
+    def load_prompt(self, prompt_file_path):
+        """Load prompt from a JSON file."""
+        with open(prompt_file_path, 'r') as file:
+            return json.load(file)
+
+    def queue_prompt(self, prompt):
+        p = {"prompt": prompt, "client_id": self.client_id}
+        data = json.dumps(p).encode('utf-8')
+        req = urllib.request.Request("http://{}/prompt".format(self.server_address), data=data)
+        return json.loads(urllib.request.urlopen(req).read())
+
+    async def get_images(self, ws, prompt):
+        prompt_id = self.queue_prompt(prompt)['prompt_id']
+        output_images = {}
+        current_node = ""
+        try:
+            while True:
+                out = await asyncio.wait_for(asyncio.to_thread(ws.recv),
+                                             timeout=60)  # Wait for a message with a 60-second timeout
+                if isinstance(out, str):
+                    message = json.loads(out)
+                    if message['type'] == 'executing':
+                        data = message['data']
+                        if data['prompt_id'] == prompt_id:
+                            if data['node'] is None:
+                                break  # Execution is done
+                            else:
+                                current_node = data['node']
+                else:
+                    if current_node == '18':
+                        logging.info(f"Received image data for prompt_id : {prompt_id}")
+                        images_output = output_images.get(current_node, [])
+                        images_output.append(out[8:])
+                        output_images[current_node] = images_output
+        except asyncio.TimeoutError:
+            logging.error("Timed out waiting for the WebSocket message.")
+            return []
+
+        return output_images
+
+    async def gen_image(self, src_image_url_1: str, src_image_url_2: str, image_name: str = "generated_image"):
+        # Establish WebSocket connection
+        ws = websocket.WebSocket()
+        ws.connect(f"ws://{self.server_address}/ws?clientId={self.client_id}")
+
+        # Generate images
+        self.prompt_data["31"]["inputs"]["image_url"] = src_image_url_1
+        self.prompt_data["32"]["inputs"]["image_url"] = src_image_url_2
+        images = await self.get_images(ws, self.prompt_data)
+
+        # Get the first image (assuming there's at least one)
+        for node_id in images:
+            for image_data in images[node_id]:
+                s3_url = upload_file_to_s3_binary(image_data, image_name)
+                return {"image_url": s3_url}
+
+        return {"error": "No images generated"}
+
+
+# Example usage
+if __name__ == "__main__":
+    server_address = "127.0.0.1:8188"
+    prompt_file_path = "path/to/your/prompt.json"
+
+    image_gen = ImageGenerator(server_address, prompt_file_path)
+    result = image_gen.gen_image(image_name="test_image")
+    print(result)

external_workers/llm_workers/models_server/blend_image_processor.py

@@ -1,12 +1,12 @@
 import os
 import random
 import sys
-from typing import Sequence, Mapping, Any, Union
+from typing import Sequence, Mapping, Any
 import torch
 import click
 
 # Functions for setup and utility
-def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+def get_value_at_index(obj, index: int) -> Any:
     """Returns the value at the given index of a sequence or mapping."""
     try:
         return obj[index]

external_workers/requirements.txt

@@ -6,4 +6,6 @@ redis==3.5.3
 
 boto3
 httpx
-openai
+openai==1.37.0
+websocket-client==1.6.1
+Pillow==10.0.0