-
Notifications
You must be signed in to change notification settings - Fork 6
/
node.py
155 lines (130 loc) · 5.52 KB
/
node.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
import os
import json
import requests
from PIL import Image
from io import BytesIO
import io
from torchvision import transforms
import torch
import base64
import time
from replicate.client import Client
from .schema_to_node import (
schema_to_comfyui_input_types,
get_return_type,
name_and_version,
inputs_that_need_arrays,
)
replicate = Client(headers={"User-Agent": "comfyui-replicate/1.0.1"})
def create_comfyui_node(schema):
replicate_model, node_name = name_and_version(schema)
return_type = get_return_type(schema)
class ReplicateToComfyUI:
@classmethod
def IS_CHANGED(cls, **kwargs):
return time.time() if kwargs["force_rerun"] else ""
@classmethod
def INPUT_TYPES(cls):
return schema_to_comfyui_input_types(schema)
RETURN_TYPES = (return_type,)
FUNCTION = "run_replicate_model"
CATEGORY = "Replicate"
def convert_input_images_to_base64(self, kwargs):
for key, value in kwargs.items():
if value is not None:
input_type = (
self.INPUT_TYPES()["required"].get(key, (None,))[0]
or self.INPUT_TYPES().get("optional", {}).get(key, (None,))[0]
)
if input_type == "IMAGE":
kwargs[key] = self.image_to_base64(value)
def image_to_base64(self, image):
if isinstance(image, torch.Tensor):
image = image.permute(0, 3, 1, 2).squeeze(0)
to_pil = transforms.ToPILImage()
pil_image = to_pil(image)
else:
pil_image = image
buffer = io.BytesIO()
pil_image.save(buffer, format="PNG")
buffer.seek(0)
img_str = base64.b64encode(buffer.getvalue()).decode()
return f"data:image/png;base64,{img_str}"
def handle_array_inputs(self, kwargs):
array_inputs = inputs_that_need_arrays(schema)
for input_name in array_inputs:
if input_name in kwargs:
if isinstance(kwargs[input_name], str):
if kwargs[input_name] == "":
kwargs[input_name] = []
else:
kwargs[input_name] = kwargs[input_name].split("\n")
else:
kwargs[input_name] = [kwargs[input_name]]
def log_input(self, kwargs):
truncated_kwargs = {
k: v[:20] + "..."
if isinstance(v, str) and v.startswith("data:image")
else v
for k, v in kwargs.items()
}
print(f"Running {replicate_model} with {truncated_kwargs}")
def handle_image_output(self, output):
# Handle both string and list outputs
output_list = [output] if isinstance(output, str) else list(output)
if output_list:
output_tensors = []
transform = transforms.ToTensor()
for image_url in output_list:
response = requests.get(image_url)
if response.status_code == 200:
image = Image.open(BytesIO(response.content))
if image.mode != "RGB":
image = image.convert("RGB")
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
return (
torch.cat(output_tensors, dim=0)
if len(output_tensors) > 1
else output_tensors[0]
)
else:
print("No output received from the model")
return None
def run_replicate_model(self, **kwargs):
self.handle_array_inputs(kwargs)
self.convert_input_images_to_base64(kwargs)
self.log_input(kwargs)
output = replicate.run(replicate_model, input=kwargs)
print(f"Output: {output}")
if return_type == "IMAGE":
output = self.handle_image_output(output)
else:
output = "".join(list(output)).strip()
return (output,)
return node_name, ReplicateToComfyUI
def create_comfyui_nodes_from_schemas(schemas_dir):
nodes = {}
current_path = os.path.dirname(os.path.abspath(__file__))
schemas_dir_path = os.path.join(current_path, schemas_dir)
for schema_file in os.listdir(schemas_dir_path):
if schema_file.endswith(".json"):
with open(os.path.join(schemas_dir_path, schema_file), "r", encoding="utf-8") as f:
schema = json.load(f)
node_name, node_class = create_comfyui_node(schema)
nodes[node_name] = node_class
return nodes
_cached_node_class_mappings = None
def get_node_class_mappings():
global _cached_node_class_mappings
if _cached_node_class_mappings is None:
_cached_node_class_mappings = create_comfyui_nodes_from_schemas("schemas")
return _cached_node_class_mappings
NODE_CLASS_MAPPINGS = get_node_class_mappings()