import time import comfy.model_management import tempfile import io from PIL import Image import requests import numpy as np import yaml import logging import folder_paths import os import sys sys.path.append(".") from comfy_api.latest._input_impl.video_types import VideoFromFile from comfy.comfy_types import IO, FileLocator, ComfyNodeABC from .utils import tensor2pil, np2tensor logger = logging.getLogger(__name__) config_dir = os.path.join(folder_paths.base_path, "config") if not os.path.exists(config_dir): os.makedirs(config_dir) def get_config(): try: config_path = os.path.join(config_dir, 'replicate_config.yml') with open(config_path, 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) return config except: return {} def save_config(config): config_path = os.path.join(config_dir, 'replicate_config.yml') with open(config_path, 'w') as f: yaml.dump(config, f, indent=4) class ComfyUIReplicateRun: """结合ComfyUI中断机制的Replicate运行器""" def __init__(self, timeout_seconds=300, check_interval=1.0): self.timeout_seconds = timeout_seconds self.check_interval = check_interval def run_with_interrupt_check(self, client, ref, input=None, **params): """带中断检查的replicate运行""" start_time = time.time() # 设置wait=False，手动控制轮询 params['wait'] = False try: # 创建预测 if hasattr(ref, 'id'): prediction = client.predictions.create( version=ref.id, input=input or {}, **params ) else: prediction = client.models.predictions.create( model=ref, input=input or {}, **params ) # 手动轮询，检查中断 while True: # 检查超时 if time.time() - start_time > self.timeout_seconds: prediction.cancel() raise Exception(f"timeout ({self.timeout_seconds} seconds)") # 检查ComfyUI中断信号 if comfy.model_management.processing_interrupted(): prediction.cancel() raise comfy.model_management.InterruptProcessingException( "ComfyUI interrupted") # 检查预测状态 prediction.reload() if prediction.status == "succeeded": return prediction.output elif prediction.status == "failed": raise Exception(f"prediction failed: {prediction.error}") elif prediction.status in ["starting", "processing"]: time.sleep(self.check_interval) else: prediction.cancel() raise Exception(f"unknown status: {prediction.status}") except Exception as e: logging.error(f"Replicate operation failed: {e}") raise class ReplicateRequstNode: def __init__(self, api_key=None): from replicate.client import Client config = get_config() self.api_key = api_key or config.get("REPLICATE_API_TOKEN") if self.api_key is not None: self.configure_replicate() self.client = Client(timeout=60) def configure_replicate(self): if self.api_key: os.environ["REPLICATE_API_TOKEN"] = self.api_key @classmethod def INPUT_TYPES(cls): return { "required": { "prompt": ("STRING", {"default": "style of 80s cyberpunk, a portrait photo", "multiline": True}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True, "tooltip": "The random seed used for creating the noise."}), "aspect_ratio": (["1:1", "16:9", "21:9", "3:2", "4:3", "5:4", "9:16", "9:21", "2:3", "3:4", "4:5"], {"default": "1:1"}), "steps": ("INT", {"default": 28, "min": 1, "max": 100}), "guidance": ("FLOAT", {"default": 3.5, "min": 0.1, "max": 100.0, "step": 0.1}), "go_fast": ("BOOLEAN", {"default": True}), }, "optional": { "api_key": ("STRING", {"default": ""}), "lora_path": ("STRING", {"default": "", "multiline": True}), "lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), "extra_lora": ("STRING", {"default": "", "multiline": True}), "extra_lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}), "model": ("STRING", {"default": "black-forest-labs/flux-dev-lora"}), "num_outputs": ("INT", {"default": 1, "min": 1, "max": 10}), "image": ("IMAGE",), "timeout": ("INT", {"default": 60, "min": 1, "max": 3000}), } } RETURN_TYPES = ("IMAGE", "INT", "INT", "STRING") RETURN_NAMES = ("image", "width", "height", "url") FUNCTION = "generate_image" CATEGORY = "utils/image" def generate_image(self, prompt, seed, aspect_ratio, steps, guidance, go_fast, lora_path="", lora_scale=1.0, api_key="", extra_lora="", extra_lora_scale=1.0, model="black-forest-labs/flux-dev-lora", num_outputs=1, image=None, timeout=60): # 更新API key if api_key.strip(): self.api_key = api_key save_config({"REPLICATE_API_TOKEN": self.api_key}) self.configure_replicate() if not self.api_key: raise ValueError( "API key not found in replicate_config.yml or node input") try: # 准备输入参数 input_params = { "prompt": prompt, "lora_weights": lora_path, "seed": seed, "aspect_ratio": aspect_ratio, "num_inference_steps": steps, "guidance": guidance, "go_fast": go_fast, "lora_scale": lora_scale, "output_format": "png", "num_outputs": num_outputs } # 添加额外的LoRA参数 if extra_lora.strip(): input_params["extra_lora"] = extra_lora input_params["extra_lora_scale"] = extra_lora_scale # 处理输入图像 if image is not None and len(image) > 0: # 将tensor转换为PIL图像，然后保存为临时文件 pil_image = tensor2pil(image[0]) # 取第一张图片 # 创建临时文件 with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file: pil_image.save(temp_file.name, format='PNG') temp_file_path = temp_file.name # 使用open()创建文件对象 input_params["input_image"] = open(temp_file_path, "rb") logger.debug(f"已添加输入图像文件: {temp_file_path}") logger.debug(f"调用Replicate API，参数: {input_params}") runner = ComfyUIReplicateRun(timeout_seconds=timeout, check_interval=1.0) # 调用Replicate API output = runner.run_with_interrupt_check(self.client, model, input=input_params) # 清理临时文件 if image is not None and len(image) > 0: try: input_params["input_image"].close() os.unlink(temp_file_path) except: pass images = [] urls = [] if not isinstance(output, list): output = [output] for image_url in output: logger.debug(f"生成的图片URL: {image_url}") urls.append(str(image_url)) response = requests.get(image_url) response.raise_for_status() image = Image.open(io.BytesIO(response.content)) width, height = image.size image_array = np.array(image) if len(image_array.shape) == 3 and image_array.shape[2] == 4: image_array = image_array[:, :, :3] images.append(image_array) image_tensor = np2tensor(images) urls_str = ",".join(urls) return (image_tensor, width, height, urls_str) except Exception as e: logger.exception(f"Replicate API调用失败: {str(e)}") raise e class ReplicateVideoRequestNode: def __init__(self, api_key=None): from replicate.client import Client config = get_config() self.api_key = api_key or config.get("REPLICATE_API_TOKEN") if self.api_key is not None: self.configure_replicate() self.client = Client(timeout=300) def configure_replicate(self): if self.api_key: os.environ["REPLICATE_API_TOKEN"] = self.api_key @classmethod def INPUT_TYPES(cls): return { "required": { "prompt": ("STRING", {"default": "", "multiline": True}), "model": ("STRING", {"default": "wan-video/wan-2.2-i2v-fast"}), "num_frames": ("INT", {"default": 81, "min": 81, "max": 121}), "resolution": (["480p", "720p"], {"default": "720p"}), "frames_per_second": ("INT", {"default": 16, "min": 5, "max": 30, "step": 1}), }, "optional": { "image": ("IMAGE", {"default": None, "tooltip": "The image for the image to video only."}), "go_fast": ("BOOLEAN", {"default": True}), "sample_shift": ("FLOAT", {"default": 12.0, "min": 1.0, "max": 20.0, "step": 0.1}), "lora_weights_transformer": ("STRING", {"default": "", "multiline": True}), "lora_scale_transformer": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}), "lora_weights_transformer_2": ("STRING", {"default": "", "multiline": True}), "lora_scale_transformer_2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}), "api_key": ("STRING", {"default": ""}), "timeout": ("INT", {"default": 300, "min": 1, "max": 3000}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True, "tooltip": "The random seed used for creating the noise."}), "disable_safety_checker": ("BOOLEAN", {"default": True}), "aspect_ratio": (["16:9", "9:16"], {"default": "16:9", "tooltip": "The aspect ratio for the text to video only."}), } } RETURN_TYPES = (IO.VIDEO, "INT", "INT", "FLOAT", "STRING") RETURN_NAMES = ("video", "width", "height", "fps", "url") FUNCTION = "generate_video" CATEGORY = "utils/video" def generate_video(self, prompt, model, num_frames, resolution, frames_per_second, image=None, go_fast=True, sample_shift=12.0, lora_weights_transformer="", lora_scale_transformer=1.0, lora_weights_transformer_2="", lora_scale_transformer_2=1.0, api_key="", timeout=300, seed=0, disable_safety_checker=True, aspect_ratio="16:9"): if api_key.strip(): self.api_key = api_key save_config({"REPLICATE_API_TOKEN": self.api_key}) self.configure_replicate() if not self.api_key: raise ValueError("API key not found in replicate_config.yml or node input") try: input_params = { "prompt": prompt, "num_frames": num_frames, "resolution": resolution, "frames_per_second": frames_per_second, "go_fast": go_fast, "sample_shift": sample_shift, "seed": seed, "disable_safety_checker": disable_safety_checker, "aspect_ratio": aspect_ratio, } temp_file_paths = [] if image is not None and len(image) > 0: for i, image_tensor in enumerate(image): pil_image = tensor2pil(image_tensor) with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file: pil_image.save(temp_file.name, format='PNG') temp_file_paths.append(temp_file.name) if i == 0: input_params["image"] = open(temp_file_paths[-1], "rb") else: input_params["last_image"] = open(temp_file_paths[-1], "rb") if lora_weights_transformer.strip(): input_params["lora_weights_transformer"] = lora_weights_transformer input_params["lora_scale_transformer"] = lora_scale_transformer if lora_weights_transformer_2.strip(): input_params["lora_weights_transformer_2"] = lora_weights_transformer_2 input_params["lora_scale_transformer_2"] = lora_scale_transformer_2 logger.debug(f"调用Replicate API生成视频，参数: {input_params}") runner = ComfyUIReplicateRun(timeout_seconds=timeout, check_interval=1.0) output = runner.run_with_interrupt_check(self.client, model, input=input_params) for i, temp_file_path in enumerate(temp_file_paths): try: if i == 0: input_params["image"].close() else: input_params["last_image"].close() os.unlink(temp_file_path) except: pass if not isinstance(output, list): output = [output] video_url = output[0] if output else None if not video_url: raise Exception("No video URL returned from API") logger.debug(f"生成的视频URL: {video_url}") videos_dir = os.path.join(folder_paths.get_output_directory(), "videos_utils_nodes") if not os.path.exists(videos_dir): os.makedirs(videos_dir) video_filename = f"replicate_video_{int(time.time())}.mp4" video_path = os.path.join(videos_dir, video_filename) response = requests.get(video_url) response.raise_for_status() with open(video_path, 'wb') as f: f.write(response.content) logger.info(f"视频已保存到: {video_path}") video_input = VideoFromFile(video_path) width, height = video_input.get_dimensions() fps = float(frames_per_second) return (video_input, width, height, fps, video_url) except Exception as e: logger.exception(f"Replicate视频生成失败: {str(e)}") raise e NODE_CLASS_MAPPINGS = { "ReplicateRequstNode": ReplicateRequstNode, "ReplicateVideoRequestNode": ReplicateVideoRequestNode, } NODE_DISPLAY_NAME_MAPPINGS = { "ReplicateVideoRequestNode": "Replicate Video Request", "ReplicateRequstNode": "Replicate Image Request", }