301 lines
10 KiB
Python
301 lines
10 KiB
Python
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"]
|