From 797180b5c72e73010b3066f876cf311eb24ee46b Mon Sep 17 00:00:00 2001 From: bennykok Date: Wed, 24 Apr 2024 21:48:42 +0800 Subject: [PATCH] feat(plugin): add external image batch --- comfy-nodes/external_image_batch.py | 85 +++++++++++++++++++++++++++++ custom_routes.py | 69 +++++++++++++---------- 2 files changed, 124 insertions(+), 30 deletions(-) create mode 100644 comfy-nodes/external_image_batch.py diff --git a/comfy-nodes/external_image_batch.py b/comfy-nodes/external_image_batch.py new file mode 100644 index 0000000..5f2c46c --- /dev/null +++ b/comfy-nodes/external_image_batch.py @@ -0,0 +1,85 @@ +import folder_paths +from PIL import Image, ImageOps +import numpy as np +import torch +import json +import comfy + +class ComfyUIDeployExternalImageBatch: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "input_id": ( + "STRING", + {"multiline": False, "default": "input_images"}, + ), + "images": ( + "STRING", + {"multiline": False, "default": "[]"}, + ), + }, + "optional": { + "default_value": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "run" + + CATEGORY = "image" + + def run(self, input_id, images=None, default_value=None): + processed_images = [] + try: + images_list = json.loads(images) # Assuming images is a JSON array string + print(images_list) + for img_input in images_list: + if img_input.startswith('http'): + import requests + from io import BytesIO + print("Fetching image from url: ", img_input) + response = requests.get(img_input) + image = Image.open(BytesIO(response.content)) + elif img_input.startswith('data:image/png;base64,') or img_input.startswith('data:image/jpeg;base64,') or img_input.startswith('data:image/jpg;base64,'): + import base64 + from io import BytesIO + print("Decoding base64 image") + base64_image = img_input[img_input.find(",")+1:] + decoded_image = base64.b64decode(base64_image) + image = Image.open(BytesIO(decoded_image)) + else: + raise ValueError("Invalid image url or base64 data provided.") + + image = ImageOps.exif_transpose(image) + image = image.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image)[None,] + processed_images.append(image_tensor) + except Exception as e: + print(f"Error processing images: {e}") + pass + + if default_value is not None and len(images_list) == 0: + processed_images.append(default_value) # Assuming default_value is a pre-processed image tensor + + # Resize images if necessary and concatenate from MakeImageBatch in ImpactPack + if processed_images: + base_shape = processed_images[0].shape[1:] # Get the shape of the first image for comparison + batch_tensor = processed_images[0] + for i in range(1, len(processed_images)): + if processed_images[i].shape[1:] != base_shape: + # Resize to match the first image's dimensions + processed_images[i] = comfy.utils.common_upscale(processed_images[i].movedim(-1, 1), base_shape[1], base_shape[0], "lanczos", "center").movedim(1, -1) + + batch_tensor = torch.cat((batch_tensor, processed_images[i]), dim=0) + # Concatenate using torch.cat + else: + batch_tensor = None # or handle the empty case as needed + return (batch_tensor, ) + + +NODE_CLASS_MAPPINGS = {"ComfyUIDeployExternalImageBatch": ComfyUIDeployExternalImageBatch} +NODE_DISPLAY_NAME_MAPPINGS = {"ComfyUIDeployExternalImageBatch": "External Image Batch (ComfyUI Deploy)"} \ No newline at end of file diff --git a/custom_routes.py b/custom_routes.py index d83910b..ff94dec 100644 --- a/custom_routes.py +++ b/custom_routes.py @@ -105,6 +105,38 @@ def apply_random_seed_to_workflow(workflow_api): workflow_api[key]['inputs']['seed'] = randomSeed(8); continue workflow_api[key]['inputs']['seed'] = randomSeed(); + +def apply_inputs_to_workflow(workflow_api: Any, inputs: Any, sid: str = None): + # Loop through each of the inputs and replace them + for key, value in workflow_api.items(): + if 'inputs' in value: + + # Support websocket + if sid is not None: + if (value["class_type"] == "ComfyDeployWebscoketImageOutput"): + value['inputs']["client_id"] = sid + if (value["class_type"] == "ComfyDeployWebscoketImageInput"): + value['inputs']["client_id"] = sid + + if "input_id" in value['inputs'] and value['inputs']['input_id'] in inputs: + new_value = inputs[value['inputs']['input_id']] + + # Lets skip it if its an image + if isinstance(new_value, Image.Image): + continue + + # Backward compactibility + value['inputs']["input_id"] = new_value + + # Fix for external text default value + if (value["class_type"] == "ComfyUIDeployExternalText"): + value['inputs']["default_value"] = new_value + + if (value["class_type"] == "ComfyUIDeployExternalCheckpoint"): + value['inputs']["default_value"] = new_value + + if (value["class_type"] == "ComfyUIDeployExternalImageBatch"): + value['inputs']["images"] = new_value def send_prompt(sid: str, inputs: StreamingPrompt): # workflow_api = inputs.workflow_api @@ -114,31 +146,8 @@ def send_prompt(sid: str, inputs: StreamingPrompt): apply_random_seed_to_workflow(workflow_api) print("getting inputs" , inputs.inputs) - - # Loop through each of the inputs and replace them - for key, value in workflow_api.items(): - if 'inputs' in value: - if (value["class_type"] == "ComfyDeployWebscoketImageOutput"): - value['inputs']["client_id"] = sid - if (value["class_type"] == "ComfyDeployWebscoketImageInput"): - value['inputs']["client_id"] = sid - - if "input_id" in value['inputs'] and value['inputs']['input_id'] in inputs.inputs: - new_value = inputs.inputs[value['inputs']['input_id']] - - # Lets skip it if its an image - if isinstance(new_value, Image.Image): - continue - - value['inputs']["input_id"] = new_value - - # Fix for external text default value - if (value["class_type"] == "ComfyUIDeployExternalText"): - value['inputs']["default_value"] = new_value - - if (value["class_type"] == "ComfyUIDeployExternalCheckpoint"): - value['inputs']["default_value"] = new_value - + + apply_inputs_to_workflow(workflow_api, inputs.inputs, sid=sid) print(workflow_api) @@ -171,15 +180,15 @@ async def comfy_deploy_run(request): prompt_server = server.PromptServer.instance data = await request.json() - workflow_api = data.get("workflow_api") - + # In older version, we use workflow_api, but this has inputs already swapped in nextjs frontend, which is tricky + workflow_api = data.get("workflow_api_raw") # The prompt id generated from comfy deploy, can be None prompt_id = data.get("prompt_id") + inputs = data.get("inputs") + # Now it handles directly in here apply_random_seed_to_workflow(workflow_api) - # for key in workflow_api: - # if 'inputs' in workflow_api[key] and 'seed' in workflow_api[key]['inputs']: - # workflow_api[key]['inputs']['seed'] = randomSeed() + apply_inputs_to_workflow(workflow_api, inputs) prompt = { "prompt": workflow_api,