import os import json from PIL import Image from io import BytesIO import io from torchvision import transforms import torch import base64 import time import torchaudio import soundfile as sf 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 = ( tuple(return_type.values()) if isinstance(return_type, dict) else (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) elif input_type == "AUDIO": kwargs[key] = self.audio_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 audio_to_base64(self, audio): if ( isinstance(audio, dict) and "waveform" in audio and "sample_rate" in audio ): waveform = audio["waveform"] sample_rate = audio["sample_rate"] else: waveform, sample_rate = audio # Ensure waveform is 2D if waveform.dim() == 1: waveform = waveform.unsqueeze(0) elif waveform.dim() > 2: waveform = waveform.squeeze() if waveform.dim() > 2: raise ValueError("Waveform must be 1D or 2D") buffer = io.BytesIO() sf.write(buffer, waveform.numpy().T, sample_rate, format="wav") buffer.seek(0) audio_str = base64.b64encode(buffer.getvalue()).decode() return f"data:audio/wav;base64,{audio_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") or v.startswith("data:audio")) else v for k, v in kwargs.items() } print(f"Running {replicate_model} with {truncated_kwargs}") def handle_image_output(self, output): if output is None: print("No image output received") return None output_list = [output] if not isinstance(output, list) else output if output_list: output_tensors = [] transform = transforms.ToTensor() for file_obj in output_list: image_data = file_obj.read() image = Image.open(BytesIO(image_data)) 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) # 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 handle_audio_output(self, output): if output is None: print("No audio output received from the model") return None output_list = [output] if not isinstance(output, list) else output audio_data = [] for audio_file in output_list: if audio_file: audio_content = BytesIO(audio_file.read()) waveform, sample_rate = torchaudio.load(audio_content) audio_data.append({ "waveform": waveform.unsqueeze(0), "sample_rate": sample_rate }) else: print("Empty audio file received") if len(audio_data) == 1: return audio_data[0] elif len(audio_data) > 0: return audio_data else: print("No valid audio files processed") return None def remove_falsey_optional_inputs(self, kwargs): optional_inputs = self.INPUT_TYPES().get("optional", {}) for key in list(kwargs.keys()): if key in optional_inputs: if isinstance(kwargs[key], torch.Tensor): continue elif not kwargs[key]: del kwargs[key] def run_replicate_model(self, **kwargs): self.handle_array_inputs(kwargs) self.remove_falsey_optional_inputs(kwargs) self.convert_input_images_to_base64(kwargs) self.log_input(kwargs) kwargs_without_force_rerun = { k: v for k, v in kwargs.items() if k != "force_rerun" } output = replicate.run(replicate_model, input=kwargs_without_force_rerun) print(f"Output: {output}") processed_outputs = [] if isinstance(return_type, dict): for prop_name, prop_type in return_type.items(): if prop_type == "IMAGE": processed_outputs.append( self.handle_image_output(output.get(prop_name)) ) elif prop_type == "AUDIO": processed_outputs.append( self.handle_audio_output(output.get(prop_name)) ) elif prop_type == "STRING": processed_outputs.append( "".join(list(output.get(prop_name, ""))).strip() ) else: if return_type == "IMAGE": processed_outputs.append(self.handle_image_output(output)) elif return_type == "AUDIO": processed_outputs.append(self.handle_audio_output(output)) else: processed_outputs.append("".join(list(output)).strip()) return tuple(processed_outputs) 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()