feat(plugin): add external image batch
This commit is contained in:
		
							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)"}
 | 
			
		||||
@ -106,30 +106,26 @@ def apply_random_seed_to_workflow(workflow_api):
 | 
			
		||||
                continue
 | 
			
		||||
            workflow_api[key]['inputs']['seed'] = randomSeed();
 | 
			
		||||
            
 | 
			
		||||
def send_prompt(sid: str, inputs: StreamingPrompt):
 | 
			
		||||
    # workflow_api = inputs.workflow_api
 | 
			
		||||
    workflow_api = copy.deepcopy(inputs.workflow_api)
 | 
			
		||||
    
 | 
			
		||||
    # Random seed
 | 
			
		||||
    apply_random_seed_to_workflow(workflow_api)
 | 
			
		||||
            
 | 
			
		||||
    print("getting inputs" , inputs.inputs)
 | 
			
		||||
            
 | 
			
		||||
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.inputs:
 | 
			
		||||
                new_value = inputs.inputs[value['inputs']['input_id']]
 | 
			
		||||
            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
 | 
			
		||||
@ -139,6 +135,19 @@ def send_prompt(sid: str, inputs: StreamingPrompt):
 | 
			
		||||
                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
 | 
			
		||||
    workflow_api = copy.deepcopy(inputs.workflow_api)
 | 
			
		||||
    
 | 
			
		||||
    # Random seed
 | 
			
		||||
    apply_random_seed_to_workflow(workflow_api)
 | 
			
		||||
            
 | 
			
		||||
    print("getting inputs" , inputs.inputs)
 | 
			
		||||
    
 | 
			
		||||
    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