diff --git a/node.py b/node.py index c706536..240f81e 100644 --- a/node.py +++ b/node.py @@ -1,15 +1,22 @@ import os import json import replicate +import requests +from PIL import Image +from io import BytesIO +from torchvision import transforms +import torch -def convert_type(openapi_type): +def convert_to_comfyui_input_type(openapi_type, openapi_format=None): type_mapping = { "string": "STRING", "integer": "INT", "number": "FLOAT", - "boolean": "BOOL", + "boolean": "BOOLEAN", } + if openapi_type == "string" and openapi_format == "uri": + return "IMAGE" return type_mapping.get(openapi_type, "STRING") @@ -24,7 +31,9 @@ def resolve_schema(prop_data, schemas): def convert_schema_to_comfyui(schema, schemas): - input_types = {"required": {}} + input_types = {"required": {}, "optional": {}} + + required_props = schema.get("required", []) for prop_name, prop_data in schema["properties"].items(): prop_data = resolve_schema(prop_data, schemas) @@ -35,11 +44,13 @@ def convert_schema_to_comfyui(schema, schemas): if "enum" in prop_data: input_type = prop_data["enum"] elif "type" in prop_data: - input_type = convert_type(prop_data["type"]) + input_type = convert_to_comfyui_input_type( + prop_data["type"], prop_data.get("format") + ) else: input_type = "STRING" - default_value = prop_data.get("default", "") + default_value = prop_data.get("default", None) input_config = {"default": default_value} @@ -48,34 +59,61 @@ def convert_schema_to_comfyui(schema, schemas): if "maximum" in prop_data: input_config["max"] = prop_data["maximum"] - if prop_data.get("type") == "string" and prop_data.get("format") == "uri": - input_config["multiline"] = True - if "prompt" in prop_name and prop_data.get("type") == "string": input_config["multiline"] = True + # Meta prompt_template needs `{prompt}` to be sent through if "template" not in prop_name: input_config["dynamicPrompts"] = True - input_types["required"][prop_name] = (input_type, input_config) + if prop_name in required_props: + input_types["required"][prop_name] = (input_type, input_config) + else: + input_types["optional"][prop_name] = (input_type, input_config) - # Reorder input_types to put prompt and negative_prompt first - ordered_input_types = {"required": {}} - for key in ["prompt", "negative_prompt"]: - if key in input_types["required"]: - ordered_input_types["required"][key] = input_types["required"][key] - for key in list(input_types["required"].keys()): - if "prompt" in key: - ordered_input_types["required"][key] = input_types["required"][key] - for key in list(input_types["required"].keys()): - if key not in ordered_input_types["required"] and key != "seed": - ordered_input_types["required"][key] = input_types["required"][key] - if "seed" in input_types["required"]: - ordered_input_types["required"]["seed"] = input_types["required"]["seed"] + return reorder_input_types(input_types, schema) + + +def reorder_input_types(input_types, schema): + ordered_input_types = {"required": {}, "optional": {}} + + # Sort properties based on "x-order" if available + sorted_properties = sorted( + schema["properties"].items(), key=lambda x: x[1].get("x-order", float("inf")) + ) + + for prop_name, _ in sorted_properties: + if prop_name in input_types["required"]: + ordered_input_types["required"][prop_name] = input_types["required"][ + prop_name + ] + elif prop_name in input_types["optional"]: + ordered_input_types["optional"][prop_name] = input_types["optional"][ + prop_name + ] return ordered_input_types +def get_return_type(schemas, model_info): + output_schema = schemas["components"]["schemas"].get("Output") + default_example_output = model_info.get("default_example", {}).get("output", []) + + if ( + output_schema + and output_schema.get("type") == "array" + and output_schema["items"].get("type") == "string" + and output_schema["items"].get("format") == "uri" + and default_example_output + and default_example_output[0] + and default_example_output[0] + .lower() + .endswith((".png", ".jpg", ".jpeg", ".gif", ".webp")) + ): + return "IMAGE" + return "STRING" + + def create_comfyui_node(schemas, model_info): author = model_info["owner"] name = model_info["name"] @@ -84,21 +122,60 @@ def create_comfyui_node(schemas, model_info): replicate_model = f"{author}/{name}:{version}" node_name = f"Replicate {author}/{name}" input_schema = schemas["components"]["schemas"]["Input"] + return_type = get_return_type(schemas, model_info) + print(f"{node_name} - {return_type}") class ReplicateToComfyUI: @classmethod def INPUT_TYPES(cls): return convert_schema_to_comfyui(input_schema, schemas) - RETURN_TYPES = ("STRING",) + RETURN_TYPES = (return_type,) FUNCTION = "run_openapi_to_comfyui" CATEGORY = "Replicate" def run_openapi_to_comfyui(self, **kwargs): print(f"Running {replicate_model} with {kwargs}") output = replicate.run(replicate_model, input=kwargs) - output = "".join(output).strip() - # print(f"Output: {output}") + print(f"Output: {output}") + + if return_type == "IMAGE": + # Convert generator to list + output_list = list(output) + if output_list: + output_tensors = [] + transform = transforms.ToTensor() + for image_url in output_list: + # Download the image from the URL + response = requests.get(image_url) + if response.status_code == 200: + image = Image.open(BytesIO(response.content)) + # Convert image to RGB if it's not already + if image.mode != "RGB": + image = image.convert("RGB") + # Convert to tensor and reshape + tensor_image = transform(image) + tensor_image = tensor_image.unsqueeze(0) + tensor_image = ( + tensor_image.permute(0, 2, 3, 1).cpu().float() + ) + output_tensors.append(tensor_image) + else: + print( + f"Failed to download image. Status code: {response.status_code}" + ) + # Combine all tensors into a single batch if multiple images + output = ( + torch.cat(output_tensors, dim=0) + if len(output_tensors) > 1 + else output_tensors[0] + ) + else: + print("No output received from the model") + output = None + else: + output = "".join(list(output)) + return (output,) return node_name, ReplicateToComfyUI @@ -119,16 +196,7 @@ def create_comfyui_nodes_from_schemas(schemas_dir): return nodes -# Create ComfyUI nodes for all schema files in the "schemas" directory -comfyui_nodes = create_comfyui_nodes_from_schemas("schemas") - -# Print the resulting node classes -for schema_file, node_class in comfyui_nodes.items(): - print(f"Node class for {schema_file}:") - print(node_class.INPUT_TYPES()) - - -NODE_CLASS_MAPPINGS = comfyui_nodes +NODE_CLASS_MAPPINGS = create_comfyui_nodes_from_schemas("schemas") print(NODE_CLASS_MAPPINGS) # # Load the schema