From 5da1a4a8c2227d88d9a6d90a203767149a70b596 Mon Sep 17 00:00:00 2001 From: nofunstudio Date: Sat, 25 Jan 2025 21:56:11 -0700 Subject: [PATCH] replicate node --- replicate_API.py | 311 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + 2 files changed, 312 insertions(+) create mode 100644 replicate_API.py create mode 100644 requirements.txt diff --git a/replicate_API.py b/replicate_API.py new file mode 100644 index 0000000..ce2b150 --- /dev/null +++ b/replicate_API.py @@ -0,0 +1,311 @@ +import os +import time +import requests +import tempfile +from PIL import Image +import numpy as np +import torch +import replicate +from io import BytesIO +import json +import threading +import asyncio + +class APIGenerateReplicate: + @classmethod + def INPUT_TYPES(cls): + """ + We add an OPTIONAL input 'image_prompt' of type IMAGE, + so the user can connect a Load Image node. + """ + return { + "required": { + "api_token": ("STRING", { + "multiline": False, + "default": "", + "display": "Replicate API Token" + }), + "model": ([ + "black-forest-labs/flux-1.1-pro-ultra", + # add any other Replicate models you want here... + ], { + "default": "black-forest-labs/flux-1.1-pro-ultra" + }), + "prompt": ("STRING", { + "multiline": True, + "default": "A blackhole in space" + }), + "number_of_images": ("INT", { + "default": 1, + "min": 1, + "max": 10, + "step": 1 + }), + "seed": ("INT", { + "default": -1, + "min": -1, + "max": 2147483647 + }), + "timeout": ("INT", { + "default": 300, + "min": 60, + "max": 1800, + "step": 60, + "display": "Timeout (seconds)" + }), + }, + "optional": { + # This allows an optional image input + "image_prompt": ("IMAGE", {}) + } + } + + RETURN_TYPES = ("IMAGE", "STRING",) + RETURN_NAMES = ("image", "generation_info",) + FUNCTION = "generate" + CATEGORY = "Replicate" + + def __init__(self): + self.output_dir = "output/API/Replicate" + self.metadata_dir = "output/API/Replicate/metadata" + os.makedirs(self.output_dir, exist_ok=True) + os.makedirs(self.metadata_dir, exist_ok=True) + self._interrupt_event = threading.Event() + + def get_next_number(self): + files = [f for f in os.listdir(self.output_dir) if f.endswith('.png')] + if not files: + return 1 + numbers = [int(f.split('.')[0]) for f in files] + return max(numbers) + 1 + + def create_filename(self, number): + return f"{number:03d}.png" + + def save_image_and_metadata(self, img, generation_info, number): + filename = self.create_filename(number) + filepath = os.path.join(self.output_dir, filename) + + # Save image + img.save(filepath) + + # Create metadata filename + metadata_filename = f"{number:03d}_metadata.json" + metadata_filepath = os.path.join(self.metadata_dir, metadata_filename) + + # Save metadata + with open(metadata_filepath, 'w', encoding='utf-8') as f: + json.dump(generation_info, f, indent=4, ensure_ascii=False) + + return filepath, metadata_filepath + + async def generate_single_image_async(self, input_data, api_token, model, image_prompt_tensor=None): + """ + A single async "job" that calls replicate.run() in a worker thread. + If 'image_prompt_tensor' is provided, convert it to a local file + and pass it as 'image_prompt'. + """ + try: + # Set the Replicate API token + os.environ["REPLICATE_API_TOKEN"] = api_token + + # If we have an image prompt from ComfyUI, turn it into a file + image_prompt_file = None + if image_prompt_tensor is not None: + # Convert the incoming torch tensor to PIL + pil_image = self.tensor_to_pil(image_prompt_tensor) + + # Write to a temporary file + tmp_fd, tmp_filename = tempfile.mkstemp(suffix=".png") + os.close(tmp_fd) # We'll reopen as "rb" below. + pil_image.save(tmp_filename, format="PNG") + + # We'll open the file for reading as replicate expects a file object + image_prompt_file = open(tmp_filename, "rb") + # Provide it in the input dictionary + input_data["image_prompt"] = image_prompt_file + + # Actually call replicate.run(...) in a thread to avoid blocking + def replicate_run_wrapper(): + return replicate.run(model, input=input_data) + + output = await asyncio.to_thread(replicate_run_wrapper) + + # Clean up the image prompt file if we used one + if image_prompt_file is not None: + image_prompt_file.close() + os.remove(image_prompt_file.name) + + if not output: + raise ValueError("No valid result received from replicate.run().") + + # Some Replicate models return a list of URLs, or a single URL/string + if isinstance(output, list): + image_url = output[0] + else: + image_url = output + + # Download the resulting image + image_response = requests.get(image_url) + if image_response.status_code != 200: + raise ConnectionError(f"Failed to download image: Status code {image_response.status_code}") + + # Convert into PIL + img = Image.open(BytesIO(image_response.content)).convert("RGB") + + # Save the image and metadata + number = self.get_next_number() + generation_info = { + "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), + "parameters": input_data, + "replicate_output": output + } + image_path, metadata_path = self.save_image_and_metadata(img, generation_info, number) + print(f"Saved image to: {image_path}") + print(f"Saved metadata to: {metadata_path}") + + # Convert to a torch tensor (Batch x Height x Width x Channels) + img_tensor = torch.from_numpy(np.array(img).astype(np.float32) / 255.0) + img_tensor = img_tensor.unsqueeze(0) # Add batch dimension + + return img_tensor, generation_info + + except Exception as e: + print(f"Generation error: {str(e)}") + raise Exception(f"Error generating image: {str(e)}") + + def generate(self, api_token, model, prompt, number_of_images=1, seed=-1, timeout=300, image_prompt=None): + """ + The main entry point for ComfyUI. Gathers multiple images if requested, + passes a (possibly optional) image_prompt, and returns combined results. + """ + if not api_token: + raise ValueError("API token is required.") + + self._interrupt_event.clear() + empty_image = torch.zeros((1, 1024, 1024, 3)) + + try: + images = [] + infos = [] + failed_jobs = [] + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + async def process_all_images(): + tasks = [] + for i in range(number_of_images): + if self._interrupt_event.is_set(): + break + + # Build input data + input_data = { + "prompt": prompt, + "raw": False, + "output_format": "png", + } + + # Seed logic + if seed != -1: + current_seed = seed + i + else: + current_seed = np.random.randint(0, 2147483647) + + input_data["seed"] = current_seed + + # Add more replicate model parameters if needed here + # e.g., "aspect_ratio", "safety_tolerance", etc. + + tasks.append( + self.generate_single_image_async( + input_data, + api_token, + model, + image_prompt_tensor=image_prompt # The optional image + ) + ) + + return await asyncio.gather(*tasks, return_exceptions=True) + + try: + results = loop.run_until_complete(process_all_images()) + finally: + loop.close() + + # Process the results + for result in results: + if isinstance(result, Exception): + failed_jobs.append({'error': str(result)}) + else: + img_tensor, generation_info = result + images.append(img_tensor) + infos.append(generation_info) + + # If no images were successful + if not images: + generation_info = { + "error": "All generation jobs failed", + "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), + "failed_jobs": failed_jobs + } + return (empty_image, json.dumps(generation_info, indent=2)) + + # Combine all images into a batch + combined_tensor = torch.cat(images, dim=0) + + combined_info = { + "successful_generations": len(images), + "total_requested": number_of_images, + "generation_parameters": { + "prompt": prompt, + "initial_seed": seed + }, + "individual_results": infos, + "failed_jobs": failed_jobs if failed_jobs else None + } + + return (combined_tensor, json.dumps(combined_info, indent=2)) + + except Exception as e: + generation_info = { + "error": f"Replicate generation failed: {str(e)}", + "timestamp": time.strftime("%Y-%m-%d %H:%M:%S") + } + return (empty_image, json.dumps(generation_info, indent=2)) + + def tensor_to_pil(self, tensor): + """ + Helper method: Convert a ComfyUI image tensor (CxHxW or HxWxC) to a PIL Image. + Typically, ComfyUI image is (B, H, W, C) or (H, W, C). Let's handle the common case. + """ + # If the tensor is 4D, remove batch dimension + if len(tensor.shape) == 4: + # shape is (B, H, W, C), typically B=1 + tensor = tensor[0] + + # Now shape should be (H, W, C) + arr = tensor.cpu().numpy() + arr = (arr * 255).clip(0, 255).astype("uint8") + pil_image = Image.fromarray(arr) + return pil_image + + @classmethod + def IS_CHANGED(cls, **kwargs): + return float("NaN") + + def interrupt(self): + print("Interrupting Replicate generation...") + self._interrupt_event.set() + +NODE_CLASS_MAPPINGS = { + "APIGenerateReplicate": APIGenerateReplicate +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "APIGenerateReplicate": "Replicate API" +} + +__all__ = ["APIGenerateReplicate"] + + \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a58c23d --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +replicate