refac

backend/open_webui/utils/images/comfyui.py

@@ -13,99 +13,97 @@ from pydantic import BaseModel
 
 log = logging.getLogger(__name__)
 
-default_headers = {"User-Agent": "Mozilla/5.0"}
+default_headers = {'User-Agent': 'Mozilla/5.0'}
 
 
 def queue_prompt(prompt, client_id, base_url, api_key):
-    log.info("queue_prompt")
-    p = {"prompt": prompt, "client_id": client_id}
-    data = json.dumps(p).encode("utf-8")
-    log.debug(f"queue_prompt data: {data}")
+    log.info('queue_prompt')
+    p = {'prompt': prompt, 'client_id': client_id}
+    data = json.dumps(p).encode('utf-8')
+    log.debug(f'queue_prompt data: {data}')
     try:
         req = urllib.request.Request(
-            f"{base_url}/prompt",
+            f'{base_url}/prompt',
             data=data,
-            headers={**default_headers, "Authorization": f"Bearer {api_key}"},
+            headers={**default_headers, 'Authorization': f'Bearer {api_key}'},
         )
         response = urllib.request.urlopen(req).read()
         return json.loads(response)
     except Exception as e:
-        log.exception(f"Error while queuing prompt: {e}")
+        log.exception(f'Error while queuing prompt: {e}')
         raise e
 
 
 def get_image(filename, subfolder, folder_type, base_url, api_key):
-    log.info("get_image")
-    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+    log.info('get_image')
+    data = {'filename': filename, 'subfolder': subfolder, 'type': folder_type}
     url_values = urllib.parse.urlencode(data)
     req = urllib.request.Request(
-        f"{base_url}/view?{url_values}",
-        headers={**default_headers, "Authorization": f"Bearer {api_key}"},
+        f'{base_url}/view?{url_values}',
+        headers={**default_headers, 'Authorization': f'Bearer {api_key}'},
     )
     with urllib.request.urlopen(req) as response:
         return response.read()
 
 
 def get_image_url(filename, subfolder, folder_type, base_url):
-    log.info("get_image")
-    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+    log.info('get_image')
+    data = {'filename': filename, 'subfolder': subfolder, 'type': folder_type}
     url_values = urllib.parse.urlencode(data)
-    return f"{base_url}/view?{url_values}"
+    return f'{base_url}/view?{url_values}'
 
 
 def get_history(prompt_id, base_url, api_key):
-    log.info("get_history")
+    log.info('get_history')
 
     req = urllib.request.Request(
-        f"{base_url}/history/{prompt_id}",
-        headers={**default_headers, "Authorization": f"Bearer {api_key}"},
+        f'{base_url}/history/{prompt_id}',
+        headers={**default_headers, 'Authorization': f'Bearer {api_key}'},
     )
     with urllib.request.urlopen(req) as response:
         return json.loads(response.read())
 
 
 def get_images(ws, workflow, client_id, base_url, api_key):
-    prompt_id = queue_prompt(workflow, client_id, base_url, api_key)["prompt_id"]
+    prompt_id = queue_prompt(workflow, client_id, base_url, api_key)['prompt_id']
     output_images = []
     while True:
         out = ws.recv()
         if isinstance(out, str):
             message = json.loads(out)
-            if message["type"] == "executing":
-                data = message["data"]
-                if data["node"] is None and data["prompt_id"] == prompt_id:
+            if message['type'] == 'executing':
+                data = message['data']
+                if data['node'] is None and data['prompt_id'] == prompt_id:
                     break  # Execution is done
         else:
             continue  # previews are binary data
 
     history = get_history(prompt_id, base_url, api_key)[prompt_id]
-    for node_id in history["outputs"]:
-        node_output = history["outputs"][node_id]
-        if node_id in workflow and workflow[node_id].get("class_type") in [
-            "SaveImage",
-            "PreviewImage",
+    for node_id in history['outputs']:
+        node_output = history['outputs'][node_id]
+        if node_id in workflow and workflow[node_id].get('class_type') in [
+            'SaveImage',
+            'PreviewImage',
         ]:
-            if "images" in node_output:
-                for image in node_output["images"]:
-                    url = get_image_url(
-                        image["filename"], image["subfolder"], image["type"], base_url
-                    )
-                    output_images.append({"url": url})
-    return {"data": output_images}
+            if 'images' in node_output:
+                for image in node_output['images']:
+                    url = get_image_url(image['filename'], image['subfolder'], image['type'], base_url)
+                    output_images.append({'url': url})
+    return {'data': output_images}
 
 
 async def comfyui_upload_image(image_file_item, base_url, api_key):
-    url = f"{base_url}/api/upload/image"
+    url = f'{base_url}/api/upload/image'
     headers = {}
 
     if api_key:
-        headers["Authorization"] = f"Bearer {api_key}"
+        headers['Authorization'] = f'Bearer {api_key}'
 
     _, (filename, file_bytes, mime_type) = image_file_item
 
     form = aiohttp.FormData()
-    form.add_field("image", file_bytes, filename=filename, content_type=mime_type)
-    form.add_field("type", "input")  # required by ComfyUI
+    form.add_field('image', file_bytes, filename=filename, content_type=mime_type)
+    form.add_field('type', 'input')  # required by ComfyUI
 
     async with aiohttp.ClientSession() as session:
         async with session.post(url, data=form, headers=headers) as resp:
@@ -116,7 +114,7 @@ async def comfyui_upload_image(image_file_item, base_url, api_key):
 class ComfyUINodeInput(BaseModel):
     type: Optional[str] = None
     node_ids: list[str] = []
-    key: Optional[str] = "text"
+    key: Optional[str] = 'text'
     value: Optional[str] = None
 
 
@@ -138,76 +136,56 @@ class ComfyUICreateImageForm(BaseModel):
     seed: Optional[int] = None
 
 
-async def comfyui_create_image(
-    model: str, payload: ComfyUICreateImageForm, client_id, base_url, api_key
-):
-    ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://")
+async def comfyui_create_image(model: str, payload: ComfyUICreateImageForm, client_id, base_url, api_key):
+    ws_url = base_url.replace('http://', 'ws://').replace('https://', 'wss://')
     workflow = json.loads(payload.workflow.workflow)
 
     for node in payload.workflow.nodes:
         if node.type:
-            if node.type == "model":
+            if node.type == 'model':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][node.key] = model
-            elif node.type == "prompt":
+                    workflow[node_id]['inputs'][node.key] = model
+            elif node.type == 'prompt':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "text"
-                    ] = payload.prompt
-            elif node.type == "negative_prompt":
+                    workflow[node_id]['inputs'][node.key if node.key else 'text'] = payload.prompt
+            elif node.type == 'negative_prompt':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "text"
-                    ] = payload.negative_prompt
-            elif node.type == "width":
+                    workflow[node_id]['inputs'][node.key if node.key else 'text'] = payload.negative_prompt
+            elif node.type == 'width':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "width"
-                    ] = payload.width
-            elif node.type == "height":
+                    workflow[node_id]['inputs'][node.key if node.key else 'width'] = payload.width
+            elif node.type == 'height':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "height"
-                    ] = payload.height
-            elif node.type == "n":
+                    workflow[node_id]['inputs'][node.key if node.key else 'height'] = payload.height
+            elif node.type == 'n':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "batch_size"
-                    ] = payload.n
-            elif node.type == "steps":
+                    workflow[node_id]['inputs'][node.key if node.key else 'batch_size'] = payload.n
+            elif node.type == 'steps':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "steps"
-                    ] = payload.steps
-            elif node.type == "seed":
-                seed = (
-                    payload.seed
-                    if payload.seed
-                    else random.randint(0, 1125899906842624)
-                )
+                    workflow[node_id]['inputs'][node.key if node.key else 'steps'] = payload.steps
+            elif node.type == 'seed':
+                seed = payload.seed if payload.seed else random.randint(0, 1125899906842624)
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][node.key] = seed
+                    workflow[node_id]['inputs'][node.key] = seed
         else:
             for node_id in node.node_ids:
-                workflow[node_id]["inputs"][node.key] = node.value
+                workflow[node_id]['inputs'][node.key] = node.value
 
     try:
         ws = websocket.WebSocket()
-        headers = {"Authorization": f"Bearer {api_key}"}
-        ws.connect(f"{ws_url}/ws?clientId={client_id}", header=headers)
-        log.info("WebSocket connection established.")
+        headers = {'Authorization': f'Bearer {api_key}'}
+        ws.connect(f'{ws_url}/ws?clientId={client_id}', header=headers)
+        log.info('WebSocket connection established.')
     except Exception as e:
-        log.exception(f"Failed to connect to WebSocket server: {e}")
+        log.exception(f'Failed to connect to WebSocket server: {e}')
         return None
 
     try:
-        log.info("Sending workflow to WebSocket server.")
-        log.info(f"Workflow: {workflow}")
-        images = await asyncio.to_thread(
-            get_images, ws, workflow, client_id, base_url, api_key
-        )
+        log.info('Sending workflow to WebSocket server.')
+        log.info(f'Workflow: {workflow}')
+        images = await asyncio.to_thread(get_images, ws, workflow, client_id, base_url, api_key)
     except Exception as e:
-        log.exception(f"Error while receiving images: {e}")
+        log.exception(f'Error while receiving images: {e}')
         images = None
 
     ws.close()
@@ -228,85 +206,65 @@ class ComfyUIEditImageForm(BaseModel):
     seed: Optional[int] = None
 
 
-async def comfyui_edit_image(
-    model: str, payload: ComfyUIEditImageForm, client_id, base_url, api_key
-):
-    ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://")
+async def comfyui_edit_image(model: str, payload: ComfyUIEditImageForm, client_id, base_url, api_key):
+    ws_url = base_url.replace('http://', 'ws://').replace('https://', 'wss://')
     workflow = json.loads(payload.workflow.workflow)
 
     for node in payload.workflow.nodes:
         if node.type:
-            if node.type == "model":
+            if node.type == 'model':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][node.key] = model
-            elif node.type == "image":
+                    workflow[node_id]['inputs'][node.key] = model
+            elif node.type == 'image':
                 if isinstance(payload.image, list):
                     # check if multiple images are provided
                     for idx, node_id in enumerate(node.node_ids):
                         if idx < len(payload.image):
-                            workflow[node_id]["inputs"][node.key] = payload.image[idx]
+                            workflow[node_id]['inputs'][node.key] = payload.image[idx]
                 else:
                     for node_id in node.node_ids:
-                        workflow[node_id]["inputs"][node.key] = payload.image
-            elif node.type == "prompt":
+                        workflow[node_id]['inputs'][node.key] = payload.image
+            elif node.type == 'prompt':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "text"
-                    ] = payload.prompt
-            elif node.type == "negative_prompt":
+                    workflow[node_id]['inputs'][node.key if node.key else 'text'] = payload.prompt
+            elif node.type == 'negative_prompt':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "text"
-                    ] = payload.negative_prompt
-            elif node.type == "width":
+                    workflow[node_id]['inputs'][node.key if node.key else 'text'] = payload.negative_prompt
+            elif node.type == 'width':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "width"
-                    ] = payload.width
-            elif node.type == "height":
+                    workflow[node_id]['inputs'][node.key if node.key else 'width'] = payload.width
+            elif node.type == 'height':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "height"
-                    ] = payload.height
-            elif node.type == "n":
+                    workflow[node_id]['inputs'][node.key if node.key else 'height'] = payload.height
+            elif node.type == 'n':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "batch_size"
-                    ] = payload.n
-            elif node.type == "steps":
+                    workflow[node_id]['inputs'][node.key if node.key else 'batch_size'] = payload.n
+            elif node.type == 'steps':
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][
-                        node.key if node.key else "steps"
-                    ] = payload.steps
-            elif node.type == "seed":
-                seed = (
-                    payload.seed
-                    if payload.seed
-                    else random.randint(0, 1125899906842624)
-                )
+                    workflow[node_id]['inputs'][node.key if node.key else 'steps'] = payload.steps
+            elif node.type == 'seed':
+                seed = payload.seed if payload.seed else random.randint(0, 1125899906842624)
                 for node_id in node.node_ids:
-                    workflow[node_id]["inputs"][node.key] = seed
+                    workflow[node_id]['inputs'][node.key] = seed
         else:
             for node_id in node.node_ids:
-                workflow[node_id]["inputs"][node.key] = node.value
+                workflow[node_id]['inputs'][node.key] = node.value
 
     try:
         ws = websocket.WebSocket()
-        headers = {"Authorization": f"Bearer {api_key}"}
-        ws.connect(f"{ws_url}/ws?clientId={client_id}", header=headers)
-        log.info("WebSocket connection established.")
+        headers = {'Authorization': f'Bearer {api_key}'}
+        ws.connect(f'{ws_url}/ws?clientId={client_id}', header=headers)
+        log.info('WebSocket connection established.')
     except Exception as e:
-        log.exception(f"Failed to connect to WebSocket server: {e}")
+        log.exception(f'Failed to connect to WebSocket server: {e}')
         return None
 
     try:
-        log.info("Sending workflow to WebSocket server.")
-        log.info(f"Workflow: {workflow}")
-        images = await asyncio.to_thread(
-            get_images, ws, workflow, client_id, base_url, api_key
-        )
+        log.info('Sending workflow to WebSocket server.')
+        log.info(f'Workflow: {workflow}')
+        images = await asyncio.to_thread(get_images, ws, workflow, client_id, base_url, api_key)
     except Exception as e:
-        log.exception(f"Error while receiving images: {e}")
+        log.exception(f'Error while receiving images: {e}')
         images = None
 
     ws.close()