add flux fill
This commit is contained in:
nofunstudio
+6
-3
@@ -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__ = [
|
||||
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user