nodes/replicate-comfyui-replicate
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Handle output images. Sort inputs by x-order. Fix required fields

Browse Source
This commit is contained in:
Paul
2024-06-26 13:25:27 +00:00
parent 50a70df57b
commit a60c7f4a18
1 changed files with 103 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
node.py
+103 -35
View File
@@ -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