replicate_Flux_Fill.py

import os
import time
import requests
import tempfile
import json
import numpy as np
import torch
from io import BytesIO
from PIL import Image
import replicate

class ReplicateAPI_flux_fill_pro:
    @classmethod
    def INPUT_TYPES(cls):
        """
        A single-image, synchronous node for `black-forest-labs/flux-fill-pro`.
        The parameters come directly from the model docs, plus
        two optional IMAGE inputs for 'mask' and 'image'.
        """
        return {
            "required": {
                "api_token": ("STRING", {
                    "multiline": False,
                    "default": "",
                    "display": "Replicate API Token"
                }),
                "prompt": ("STRING", {
                    "multiline": True,
                    "default": "movie poster"
                }),
                "steps": ("INT", {
                    "default": 50,
                    "min": 1,
                    "max": 200,
                    "step": 1,
                    "display": "Steps"
                }),
                "guidance": ("FLOAT", {
                    "default": 3.0,
                    "min": 0,
                    "max": 20,
                    "step": 0.1,
                    "display": "Guidance"
                }),
                "outpaint": ("STRING", {
                    "default": "Zoom out 2x",
                    "display": "Outpaint"
                }),
                "output_format": (["png", "jpg"], {
                    "default": "png",
                    "display": "Output Format"
                }),
                "safety_tolerance": ("INT", {
                    "default": 6,
                    "min": 0,
                    "max": 10,
                    "display": "Safety Tolerance"
                }),
              
            },
            "optional": {
                # Optional IMAGE inputs from ComfyUI (e.g. Load Image node)
                "mask": ("IMAGE", {
                    "display": "Mask Image (Optional)"
                }),
                "image": ("IMAGE", {
                    "display": "Source Image (Optional)"
                }),
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING",)
    RETURN_NAMES = ("image", "generation_info",)
    FUNCTION = "generate"
    CATEGORY = "Replicate"

    def __init__(self):
        # Same directory structure as your flux ultra node:
        self.output_dir = "output/API/Replicate/flux-fill-pro"
        self.metadata_dir = os.path.join(self.output_dir, "metadata")

        os.makedirs(self.output_dir, exist_ok=True)
        os.makedirs(self.metadata_dir, exist_ok=True)

    def get_next_number(self):
        """
        Looks at existing .png/.jpg files and picks the next integer
        filename index (e.g. 001, 002, etc.).
        """
        valid_exts = {".png", ".jpg"}
        files = [f for f in os.listdir(self.output_dir) 
                 if os.path.splitext(f)[1].lower() in valid_exts]
        if not files:
            return 1

        numbers = []
        for file_name in files:
            base, _ = os.path.splitext(file_name)
            try:
                numbers.append(int(base))
            except ValueError:
                pass
        
        if numbers:
            return max(numbers) + 1
        else:
            return 1

    def create_filename(self, number):
        """
        The flux ultra node used zero-padded filenames ending in .png
        (e.g., '001.png'). We'll keep that exact style here.
        """
        return f"{number:03d}.png"

    def save_image_and_metadata(self, img, generation_info, number):
        """
        Saves the image as .png, plus metadata as a .json.
        Identical to the flux ultra approach.
        """
        filename = self.create_filename(number)
        filepath = os.path.join(self.output_dir, filename)

        # Save image
        img.save(filepath, format="PNG")

        # Create metadata filename (001_metadata.json)
        metadata_filename = f"{number:03d}_metadata.json"
        metadata_filepath = os.path.join(self.metadata_dir, metadata_filename)

        # Write out 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

    def generate(
        self,
        api_token,
        prompt,
        steps,
        guidance,
        outpaint,
        output_format,
        safety_tolerance,
        mask=None,
        image=None
    ):
        """
        Single synchronous call to replicate.run("black-forest-labs/flux-fill-pro").
        Saves image & metadata exactly as done previously, then returns (IMAGE, STRING).
        """
        # For errors, we return an empty 1024x1024 (B=1,H=1024,W=1024,C=3), 
        # which is what your flux ultra node used as a fallback.
        empty_image = torch.zeros((1, 1024, 1024, 3))

        # Make sure we have an API token
        if not api_token:
            raise ValueError("A Replicate API token is required.")

        try:
            # 1) Build the input dictionary
            os.environ["REPLICATE_API_TOKEN"] = api_token
            input_data = {
                "prompt": prompt,
                "steps": steps,
                "guidance": guidance,
                "outpaint": outpaint,
                "output_format": output_format,  # "png" or "jpg"
                "safety_tolerance": safety_tolerance,
                "prompt_upsampling": True,
            }

            # 2) Convert optional mask/image to local files if present
            mask_file = None
            image_file = None

            if mask is not None:
                mask_file = self.tensor_to_tempfile(mask)
                input_data["mask"] = mask_file
            if image is not None:
                image_file = self.tensor_to_tempfile(image)
                input_data["image"] = image_file

            # 3) Call replicate.run() 
            output = replicate.run("black-forest-labs/flux-fill-pro", input=input_data)

            # 4) Clean up local files
            if mask_file is not None:
                mask_file.close()
                os.remove(mask_file.name)
            if image_file is not None:
                image_file.close()
                os.remove(image_file.name)

            if not output:
                raise ValueError("No valid result from replicate.run().")

            # 5) Usually a single URL or a list
            if isinstance(output, list):
                image_url = output[0]
            else:
                image_url = output

            # 6) Download final image
            resp = requests.get(image_url)
            if resp.status_code != 200:
                raise ConnectionError(f"Failed to download image. HTTP {resp.status_code}")

            pil_img = Image.open(BytesIO(resp.content))
            if pil_img.mode != "RGB":
                pil_img = pil_img.convert("RGB")

            # 7) Save the image & metadata with the same flux ultra logic
            number = self.get_next_number()
            # We'll remove non-serializable items from input_data
            safe_input_data = dict(input_data)
            safe_input_data.pop("mask", None)
            safe_input_data.pop("image", None)

            # Keep replicate output as string
            generation_info = {
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                "parameters": safe_input_data,
                "replicate_output": str(output),
                "model": "black-forest-labs/flux-fill-pro"
            }

            image_path, metadata_path = self.save_image_and_metadata(pil_img, generation_info, number)
            print(f"[flux-fill-pro single] Saved image -> {image_path}")
            print(f"[flux-fill-pro single] Saved metadata -> {metadata_path}")

            # 8) Convert to ComfyUI's IMAGE format (1, H, W, 3)
            img_tensor = torch.from_numpy(np.array(pil_img).astype(np.float32) / 255.0)
            img_tensor = img_tensor.unsqueeze(0)

            # 9) Return the image tensor & metadata JSON string
            return (img_tensor, json.dumps(generation_info, indent=2))

        except Exception as e:
            # Return an empty image and an error message in JSON
            error_info = {
                "error": f"Flux-fill-pro single generation failed: {str(e)}",
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
            }
            return (empty_image, json.dumps(error_info, indent=2))

    def tensor_to_tempfile(self, tensor):
        """
        Matches the flux ultra logic:
        - Convert a ComfyUI IMAGE tensor to a PNG file 
        - Return an open file in "rb" mode (caller must close/delete).
        """
        pil_img = self.tensor_to_pil(tensor)
        fd, filename = tempfile.mkstemp(suffix=".png")
        os.close(fd)
        pil_img.save(filename, format="PNG")
        return open(filename, "rb")

    def tensor_to_pil(self, tensor):
        """
        Identical to your flux ultra approach: handle (B, H, W, C) or (C, H, W).
        """
        if len(tensor.shape) == 4:
            tensor = tensor[0]  # remove batch dimension

        arr = tensor.cpu().numpy()
        # If shape is (C, H, W), transpose it
        if arr.ndim == 3 and arr.shape[0] <= 4:
            arr = np.transpose(arr, (1, 2, 0))

        arr = (arr * 255).clip(0, 255).astype("uint8")
        return Image.fromarray(arr)

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        """
        Same as your flux ultra node: 
        returning NaN means no caching for this node.
        """
        return float("NaN")

    def interrupt(self):
        """
        In single-image synchronous logic, there's no partial progress to interrupt. 
        We'll leave it here for consistency with the flux ultra node.
        """
        print("[flux-fill-pro single] Interrupt called (not used in single-image mode).")


# Register with ComfyUI
NODE_CLASS_MAPPINGS = {
    "ReplicateAPI_flux_fill_pro": ReplicateAPI_flux_fill_pro
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ReplicateAPI_flux_fill_pro": "Replicate Flux-Fill-Pro"
}

__all__ = ["ReplicateAPI_flux_fill_pro"]