182 lines
7.3 KiB
Python
182 lines
7.3 KiB
Python
import time
|
|
from .wavespeed_api.utils import imageurl2tensor
|
|
from .wavespeed_api.client import WaveSpeedClient
|
|
|
|
|
|
class SFWaveSpeedSeedreamV4EditSequential:
|
|
# Recommended resolution presets for ByteDance Seedream V4
|
|
RECOMMENDED_PRESETS_SEEDREAM_4 = [
|
|
("2048x2048 (1:1)", 2048, 2048),
|
|
("2304x1728 (4:3)", 2304, 1728),
|
|
("1728x2304 (3:4)", 1728, 2304),
|
|
("2560x1440 (16:9)", 2560, 1440),
|
|
("1440x2560 (9:16)", 1440, 2560),
|
|
("2496x1664 (3:2)", 2496, 1664),
|
|
("1664x2496 (2:3)", 1664, 2496),
|
|
("3024x1296 (21:9)", 3024, 1296),
|
|
("4096x4096 (1:1)", 4096, 4096),
|
|
]
|
|
|
|
SIZE_PRESETS = [preset[0] for preset in RECOMMENDED_PRESETS_SEEDREAM_4]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"client": ("WAVESPEED_AI_API_CLIENT",),
|
|
"prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Description of the image editing/generation task. Specify the number of images to generate and use phrases like 'a series of' or 'group of images' for consistency."
|
|
}),
|
|
"max_images": ("INT", {
|
|
"default": 2,
|
|
"min": 1,
|
|
"max": 15,
|
|
"tooltip": "Number of images to generate. Must align with prompt description."
|
|
}),
|
|
"size_preset": (s.SIZE_PRESETS, {
|
|
"default": "2048x2048 (1:1)",
|
|
"tooltip": "Resolution preset for the generated images"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff,
|
|
"control_after_generate": True,
|
|
"tooltip": "Random seed for reproducible results"
|
|
}),
|
|
"enable_sync_mode": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Wait for generation to complete before returning"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"image_1": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "First input image URL (connect from Upload Image node)"
|
|
}),
|
|
"image_2": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Second input image URL (connect from Upload Image node)"
|
|
}),
|
|
"image_3": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Third input image URL (connect from Upload Image node)"
|
|
}),
|
|
"image_4": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Fourth input image URL (connect from Upload Image node)"
|
|
}),
|
|
"image_5": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Fifth input image URL (connect from Upload Image node)"
|
|
}),
|
|
"image_6": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Sixth input image URL (connect from Upload Image node)"
|
|
}),
|
|
"image_7": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Seventh input image URL (connect from Upload Image node)"
|
|
}),
|
|
"image_8": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Eighth input image URL (connect from Upload Image node)"
|
|
}),
|
|
"image_9": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Ninth input image URL (connect from Upload Image node)"
|
|
}),
|
|
"image_10": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Tenth input image URL (connect from Upload Image node)"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("output_images",)
|
|
CATEGORY = "Stillfront/WaveSpeed"
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, client, prompt, max_images, size_preset, seed, enable_sync_mode,
|
|
image_1="", image_2="", image_3="", image_4="", image_5="",
|
|
image_6="", image_7="", image_8="", image_9="", image_10=""):
|
|
real_client = WaveSpeedClient(api_key=client["api_key"])
|
|
|
|
# Collect all provided image URLs
|
|
image_inputs = [image_1, image_2, image_3, image_4, image_5, image_6, image_7, image_8, image_9, image_10]
|
|
image_urls = [url.strip() for url in image_inputs if url and url.strip()]
|
|
|
|
# Validate max 10 images as per API documentation
|
|
if len(image_urls) > 10:
|
|
raise Exception("Maximum 10 input images allowed")
|
|
|
|
if not image_urls:
|
|
raise Exception("At least one input image is required")
|
|
|
|
# Find the preset dimensions
|
|
preset_data = next((preset for preset in self.RECOMMENDED_PRESETS_SEEDREAM_4 if preset[0] == size_preset), None)
|
|
if preset_data:
|
|
final_size = f"{preset_data[1]}*{preset_data[2]}"
|
|
else:
|
|
raise ValueError(f"Invalid size preset: {size_preset}")
|
|
|
|
payload = {
|
|
"enable_base64_output": False,
|
|
"enable_sync_mode": enable_sync_mode,
|
|
"images": image_urls,
|
|
"max_images": max_images,
|
|
"prompt": prompt,
|
|
"size": final_size
|
|
}
|
|
|
|
# Add seed if not 0 (0 means random for this API)
|
|
if seed != 0:
|
|
payload["seed"] = seed
|
|
|
|
# API endpoint for edit sequential
|
|
endpoint = "/api/v3/bytedance/seedream-v4/edit-sequential"
|
|
|
|
try:
|
|
response = real_client.post(endpoint, payload, timeout=real_client.once_timeout)
|
|
|
|
if enable_sync_mode:
|
|
# In sync mode, we get the results directly
|
|
if "outputs" in response and response["outputs"]:
|
|
image_urls = response["outputs"] # Already a list
|
|
return (imageurl2tensor(image_urls),)
|
|
else:
|
|
raise Exception(f"No output received from sync API. Response: {response}")
|
|
else:
|
|
# In async mode, we need to poll for results
|
|
task_id = response["id"]
|
|
print(f"Task submitted successfully. Request ID: {task_id}")
|
|
|
|
try:
|
|
# Wait for task to complete - longer timeout for sequential editing
|
|
result = real_client.wait_for_task(task_id, polling_interval=1, timeout=600)
|
|
|
|
if "outputs" in result and result["outputs"]:
|
|
image_urls = result["outputs"] # Already a list
|
|
return (imageurl2tensor(image_urls),)
|
|
else:
|
|
raise Exception("Task completed but no output received")
|
|
|
|
except Exception as e:
|
|
raise Exception(f"Async task failed: {str(e)}")
|
|
|
|
except Exception as e:
|
|
print(f"Error in {self.__class__.__name__}: {str(e)}")
|
|
raise e
|
|
|
|
|
|
# Node registration - REQUIRED
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SFWaveSpeedSeedreamV4EditSequential": SFWaveSpeedSeedreamV4EditSequential
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SFWaveSpeedSeedreamV4EditSequential": "SF WaveSpeed Seedream V4 Edit Sequential"
|
|
} |