feat(plugin): add external image batch
This commit is contained in:
bennykok
parent
d00ca375a2
commit
797180b5c7
85
comfy-nodes/external_image_batch.py
Normal file
85
comfy-nodes/external_image_batch.py
Normal file
@ -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)"}
|
||||
@ -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,
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user