replicate_Flux_Pro_Ultra.py

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
import io

class ReplicateAPI_flux_1_1_pro_ultra:
    @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"
                }),
                "image_prompt_strength": ("FLOAT", {
                    "default": .1,
                    "min": .1,
                    "max": 1,
                    "step": .01
                }),
                "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, image_prompt_strength=.1):
        """
        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")

            # --- Remove non-serializable file object before saving metadata ---
            safe_input_data = dict(input_data)
            # pop() ensures we remove the open file or anything else that won't serialize
            safe_input_data.pop("image_prompt", None)

            # Save the image and metadata
            number = self.get_next_number()
            generation_info = {
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                "parameters": safe_input_data,
                "replicate_output": str(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, image_prompt_strength=.1):
        """
        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, 
                            image_prompt_strength=image_prompt_strength
                        )
                    )
                
                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,
                    "image_prompt_strength": image_prompt_strength
                },
                "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 = {
    "Replicate flux 1.1 pro ultra": ReplicateAPI_flux_1_1_pro_ultra
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ReplicateAPI_flux_1_1_pro_ultra": "Replicate flux 1.1 pro ultra"
}

__all__ = ["ReplicateAPI_flux_1_1_pro_ultra"]