326 lines
12 KiB
Python
326 lines
12 KiB
Python
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"]
|