diff --git a/__init__.py b/__init__.py index 5ddc82a..ca9e51f 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,14 @@ -from .nano_seed import NanoSeedEdit - -NODE_CLASS_MAPPINGS = { - "NanoSeedEdit": NanoSeedEdit, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "NanoSeedEdit": "NanoSeed" -} - -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file +from .nano_seed import NanoSeedEdit +from .seedream_text import SeedreamTextToImage + +NODE_CLASS_MAPPINGS = { + "NanoSeedEdit": NanoSeedEdit, + "SeedreamTextToImage": SeedreamTextToImage, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NanoSeedEdit": "NanoSeed", + "SeedreamTextToImage": "Seedream Text to Image", +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/seedream_text.py b/seedream_text.py new file mode 100644 index 0000000..1e2546c --- /dev/null +++ b/seedream_text.py @@ -0,0 +1,157 @@ +import torch +import requests +import base64 +import numpy as np +from io import BytesIO +from PIL import Image + + +# Helper function to convert PIL Image (RGB/RGBA/etc.) back to ComfyUI tensor (B=1, H, W, C) +def pil2tensor(pil_image): + if pil_image is None: + return None + + if pil_image.mode != "RGB": + pil_image = pil_image.convert("RGB") + + arr = np.array(pil_image).astype(np.float32) / 255.0 + arr = arr[np.newaxis, ...] + return torch.from_numpy(arr) + + +class SeedreamTextToImage: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"default": "A beautiful cinematic image", "multiline": True}), + "model": (["seedream_5_lite", "seedream_4.5"],), + "fal_key": ("STRING", {"default": "your_fal_key_here"}), + }, + "optional": { + "image_size": ([ + "auto_2K", + "auto_3K", + "auto_4K", + "square_hd", + "square", + "portrait_4_3", + "portrait_16_9", + "landscape_4_3", + "landscape_16_9", + "custom", + ], {"default": "auto_2K"}), + "width": ("INT", {"default": 0, "min": 0, "max": 4096, "display": "number"}), + "height": ("INT", {"default": 0, "min": 0, "max": 4096, "display": "number"}), + "num_images": ("INT", {"default": 1, "min": 1, "max": 6}), + "max_images": ("INT", {"default": 1, "min": 1, "max": 6}), + "seed": ("INT", {"default": 0, "min": 0, "max": 2**32 - 1}), + "enable_safety_checker": ("BOOLEAN", {"default": True}), + "return_byteplus_urls": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("generated_image",) + FUNCTION = "generate_image" + CATEGORY = "image/generate" + OUTPUT_NODE = True + + def generate_image( + self, + prompt, + model, + fal_key, + image_size="auto_2K", + width=0, + height=0, + num_images=1, + max_images=1, + seed=0, + enable_safety_checker=True, + return_byteplus_urls=False, + ): + if fal_key == "your_fal_key_here" or not fal_key.strip(): + raise ValueError("Please set your fal.ai API key in the node.") + + if not prompt or not prompt.strip(): + raise ValueError("Prompt cannot be empty.") + + if model == "seedream_5_lite": + url = "https://fal.run/fal-ai/bytedance/seedream/v5/lite/text-to-image" + elif model == "seedream_4.5": + url = "https://fal.run/fal-ai/bytedance/seedream/v4.5/text-to-image" + else: + raise ValueError(f"Unsupported model: {model}") + + # Seedream 4.5 schema does not list auto_3K, so fall back safely. + if model == "seedream_4.5" and image_size == "auto_3K": + image_size = "auto_2K" + + payload = { + "prompt": prompt, + "num_images": min(num_images, 6), + "max_images": min(max_images, 6), + "enable_safety_checker": bool(enable_safety_checker), + "sync_mode": True, + } + + # Use custom dimensions only when both width and height are provided. + # Otherwise use fal's image_size enum. + if image_size == "custom" or (width > 0 and height > 0): + if width <= 0 or height <= 0: + raise ValueError("For custom image_size, both width and height must be greater than 0.") + if width > 4096 or height > 4096: + raise ValueError("Seedream text-to-image: width and height must be <= 4096px.") + payload["image_size"] = {"width": width, "height": height} + else: + payload["image_size"] = image_size + + # Seedream 4.5 supports seed in the public schema. V5 Lite docs currently do not list seed as an input. + if model == "seedream_4.5" and seed > 0: + payload["seed"] = seed + + # V5 Lite schema exposes this option. Keeping it conditional avoids sending unknown params to 4.5. + if model == "seedream_5_lite": + payload["return_byteplus_urls"] = bool(return_byteplus_urls) + + headers = { + "Authorization": f"Key {fal_key}", + "Content-Type": "application/json", + } + + response = requests.post(url, json=payload, headers=headers, timeout=300) + if response.status_code != 200: + raise ValueError(f"API error {response.status_code}: {response.text}") + + api_result = response.json() + if "images" not in api_result or len(api_result["images"]) == 0: + raise ValueError(f"No images returned from API. Response: {api_result}") + + output_tensors = [] + max_outputs = min(len(api_result["images"]), max(1, num_images * max_images)) + + for img_info in api_result["images"][:max_outputs]: + img_data = img_info.get("data_uri") or img_info.get("url") + if not img_data: + continue + + if img_data.startswith("data:"): + _, encoded = img_data.split(",", 1) + pil_image = Image.open(BytesIO(base64.b64decode(encoded))) + else: + img_resp = requests.get(img_data, timeout=120) + if img_resp.status_code != 200: + raise ValueError(f"Failed to download generated image: {img_resp.status_code}") + pil_image = Image.open(BytesIO(img_resp.content)) + + tensor_image = pil2tensor(pil_image) + if tensor_image is not None: + output_tensors.append(tensor_image) + + if output_tensors: + batched_output = torch.cat(output_tensors, dim=0) + else: + batched_output = torch.zeros((1, 512, 512, 3)) + + return (batched_output,)