From ed2ba571166358e728b59edc8d5d71f80ff661db Mon Sep 17 00:00:00 2001 From: nofunstudio Date: Sat, 25 Jan 2025 23:53:00 -0700 Subject: [PATCH] add flux fill --- __init__.py | 9 +- replicate_Flux_Fill.py | 312 ++++++++++++++++++ ...cate_API.py => replicate_Flux_Pro_Ultra.py | 4 +- 3 files changed, 320 insertions(+), 5 deletions(-) create mode 100644 replicate_Flux_Fill.py rename replicate_API.py => replicate_Flux_Pro_Ultra.py (98%) diff --git a/__init__.py b/__init__.py index 010a2f3..dabfaab 100644 --- a/__init__.py +++ b/__init__.py @@ -1,14 +1,17 @@ # __init__.py from .layered_infinite_zoom import LayeredInfiniteZoom -from .replicate_API import ReplicateAPI_flux_1_1_pro_ultra +from .replicate_Flux_Pro_Ultra import ReplicateAPI_flux_1_1_pro_ultra +from .replicate_Flux_Fill import ReplicateAPI_flux_fill_pro NODE_CLASS_MAPPINGS = { "LayeredInfiniteZoom": LayeredInfiniteZoom, - "ReplicateAPI_flux_1_1_pro_ultra": ReplicateAPI_flux_1_1_pro_ultra + "ReplicateAPI_flux_1_1_pro_ultra": ReplicateAPI_flux_1_1_pro_ultra, + "ReplicateAPI_flux_fill_pro": ReplicateAPI_flux_fill_pro } NODE_DISPLAY_NAME_MAPPINGS = { "LayeredInfiniteZoom": "Layered Infinite Zoom", - "ReplicateAPI_flux_1_1_pro_ultra": "Replicate flux 1.1 pro ultra" + "ReplicateAPI_flux_1_1_pro_ultra": "Replicate flux 1.1 pro ultra", + "ReplicateAPI_flux_fill_pro": "Replicate flux fill pro" } __all__ = [ diff --git a/replicate_Flux_Fill.py b/replicate_Flux_Fill.py new file mode 100644 index 0000000..fd27317 --- /dev/null +++ b/replicate_Flux_Fill.py @@ -0,0 +1,312 @@ +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 synchronous generation node for `black-forest-labs/flux-fill-pro`. + + Required fields: + - api_token (STRING) + - prompt (STRING) + - steps (INT) + - guidance (FLOAT) + - outpaint (STRING) + - output_format (choice of "jpg" or "png") + - safety_tolerance (INT) + - prompt_upsampling (BOOL) + + Optional fields (ComfyUI image inputs): + - mask (IMAGE) + - image (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 Type" + }), + "output_format": (["jpg", "png"], { + "default": "jpg", + "display": "Output Format" + }), + "safety_tolerance": ("INT", { + "default": 6, + "min": 0, + "max": 10, + "display": "Safety Tolerance" + }), + "prompt_upsampling": ("BOOL", { + "default": True, + "display": "Prompt Upsampling?" + }), + }, + "optional": { + "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): + # Output directories for saving images & metadata + self.output_dir = "output/API/Replicate/flux-fill-pro-single" + 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): + """ + Find the next available integer filename (e.g. 001.jpg). + """ + files = [f for f in os.listdir(self.output_dir) if f.endswith('.png') or f.endswith('.jpg')] + if not files: + return 1 + + numbers = [] + for f in files: + base, ext = os.path.splitext(f) + try: + numbers.append(int(base)) + except ValueError: + pass + if numbers: + return max(numbers) + 1 + else: + return 1 + + def create_filename(self, number, ext="jpg"): + """ + Format the filename like 001.jpg or 002.png, etc. + """ + return f"{number:03d}.{ext}" + + def save_image_and_metadata(self, pil_img, generation_info, number, ext="jpg"): + """ + Saves the PIL image and the metadata JSON file. Returns file paths. + """ + file_basename = self.create_filename(number, ext) + image_path = os.path.join(self.output_dir, file_basename) + + # Save the image + pil_img.save(image_path, format=ext.upper()) + + # Save JSON metadata + meta_filename = f"{number:03d}_metadata.json" + meta_path = os.path.join(self.metadata_dir, meta_filename) + with open(meta_path, 'w', encoding='utf-8') as f: + json.dump(generation_info, f, indent=4, ensure_ascii=False) + + return image_path, meta_path + + def generate( + self, + api_token, + prompt, + steps, + guidance, + outpaint, + output_format, + safety_tolerance, + prompt_upsampling, + mask=None, + image=None + ): + """ + Synchronous single-image generation for flux-fill-pro. + No multi-image or seeds. Just a single replicate.run() call. + + ComfyUI calls this function. We: + + 1) Convert 'mask' and 'image' from ComfyUI tensors to temp files (if provided). + 2) Call replicate.run("black-forest-labs/flux-fill-pro", ...) + 3) Download the result + 4) Save image & metadata + 5) Return (IMAGE, STRING) + """ + # Validate + if not api_token: + raise ValueError("API token is required for Replicate calls.") + + # We'll return a placeholder if generation fails + empty_img = torch.zeros((1, 512, 512, 3)) + + try: + # 1) Build the input dictionary + os.environ["REPLICATE_API_TOKEN"] = api_token + + input_dict = { + "prompt": prompt, + "steps": steps, + "guidance": guidance, + "outpaint": outpaint, + "output_format": output_format, + "safety_tolerance": safety_tolerance, + "prompt_upsampling": prompt_upsampling, + } + + # 2) Convert mask/image to local files if given + mask_file = None + image_file = None + + if mask is not None: + mask_file = self.tensor_to_tempfile(mask, suffix=".png") + input_dict["mask"] = mask_file + + if image is not None: + image_file = self.tensor_to_tempfile(image, suffix=".png") + input_dict["image"] = image_file + + # 3) Call replicate.run (sync) + output = replicate.run("black-forest-labs/flux-fill-pro", input=input_dict) + + # Clean up temp 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().") + + # 4) Usually a single URL or a list of URLs + if isinstance(output, list): + image_url = output[0] + else: + image_url = output + + # Download the image + resp = requests.get(image_url) + if resp.status_code != 200: + raise ConnectionError(f"Failed to download result image. Status code: {resp.status_code}") + + pil_img = Image.open(BytesIO(resp.content)).convert("RGB") + + # 5) Prepare metadata & save + number = self.get_next_number() + + # We'll remove any non-serializable items from input_dict + safe_dict = dict(input_dict) + safe_dict.pop("mask", None) + safe_dict.pop("image", None) + + generation_info = { + "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), + "model": "black-forest-labs/flux-fill-pro", + "parameters": safe_dict, + "replicate_output": str(output), # cast output to string + } + + image_path, meta_path = self.save_image_and_metadata(pil_img, generation_info, number, ext=output_format) + print(f"[flux-fill-pro] Single generation: saved image -> {image_path}") + print(f"[flux-fill-pro] Single generation: saved metadata -> {meta_path}") + + # Convert to ComfyUI torch tensor + result_tensor = torch.from_numpy(np.array(pil_img).astype(np.float32) / 255.0) + result_tensor = result_tensor.unsqueeze(0) # (1, H, W, C) + + # Return the single (IMAGE, STRING) + return (result_tensor, json.dumps(generation_info, indent=2)) + + except Exception as e: + error_info = { + "error": f"Flux-fill-pro single generation failed: {str(e)}", + "timestamp": time.strftime("%Y-%m-%d %H:%M:%S") + } + return (empty_img, json.dumps(error_info, indent=2)) + + def tensor_to_tempfile(self, tensor, suffix=".png"): + """ + Convert a ComfyUI image tensor to a local file object, so we can pass it to replicate. + We'll return an open file in "rb" mode, and the caller is responsible for closing/deleting it. + """ + pil_img = self.tensor_to_pil(tensor) + fd, filename = tempfile.mkstemp(suffix=suffix) + os.close(fd) + pil_img.save(filename, format=suffix.upper().lstrip(".")) + return open(filename, "rb") + + def tensor_to_pil(self, tensor): + """ + Convert (B, H, W, C) or (H, W, C) or (C, H, W) to a PIL image. + Typically, ComfyUI uses (B, H, W, C). We'll handle that. + """ + if len(tensor.shape) == 4: + # drop batch dimension + tensor = tensor[0] + + arr = tensor.cpu().numpy() + # If shape is (C, H, W), we transpose + if arr.ndim == 3 and arr.shape[0] <= 4: + arr = np.transpose(arr, (1, 2, 0)) + + arr = (arr * 255).clip(0, 255).astype(np.uint8) + return Image.fromarray(arr) + + @classmethod + def IS_CHANGED(cls, **kwargs): + """ + ComfyUI-specific method to handle caching logic. + Returning NaN tells ComfyUI there's no caching for this node. + """ + return float("NaN") + + def interrupt(self): + """ + For multi-step processes, we could handle an interrupt event, + but here it's just a single call. + We'll keep it for completeness. + """ + print("[flux-fill-pro-single] Interrupt: not used in single-image mode.") + +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"] diff --git a/replicate_API.py b/replicate_Flux_Pro_Ultra.py similarity index 98% rename from replicate_API.py rename to replicate_Flux_Pro_Ultra.py index 258b19f..cc928a3 100644 --- a/replicate_API.py +++ b/replicate_Flux_Pro_Ultra.py @@ -36,7 +36,7 @@ class ReplicateAPI_flux_1_1_pro_ultra: "multiline": True, "default": "A blackhole in space" }), - "image_prompt_strength": ("INT", { + "image_prompt_strength": ("FLOAT", { "default": .1, "min": .1, "max": 1, @@ -106,7 +106,7 @@ class ReplicateAPI_flux_1_1_pro_ultra: return filepath, metadata_filepath - async def generate_single_image_async(self, input_data, api_token, model, image_prompt_tensor=None): + 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