diff --git a/__init__.py b/__init__.py index 7d0f56e..5ddc82a 100644 --- a/__init__.py +++ b/__init__.py @@ -5,7 +5,7 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { - "NanoSeedEdit": "Edit image NanoBanana + Seededit by ComRender" + "NanoSeedEdit": "NanoSeed" } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/nano_seed.py b/nano_seed.py index c8fc87e..a9a464e 100644 --- a/nano_seed.py +++ b/nano_seed.py @@ -7,18 +7,18 @@ from PIL import Image # Helper function to convert ComfyUI tensor (B=1, H, W, C) to PIL Image (RGB) def tensor2pil(image_tensor): - # image_tensor is always (1, H, W, C) from the loop + if image_tensor is None or image_tensor.shape[0] == 0: + return None i = 255. * image_tensor[0].cpu().numpy() # (H, W, C) image = np.clip(i, 0, 255).astype(np.uint8) - # Handle channels for PIL (always output RGB) c = image.shape[-1] if c == 1: - image = np.repeat(image, 3, axis=-1) # (H, W, 3) + image = np.repeat(image, 3, axis=-1) elif c == 3: - pass # Already good + pass elif c == 4: - image = image[..., :3] # Drop alpha + image = image[..., :3] else: raise ValueError(f"Unsupported channels: {c}. Expected 1, 3, or 4.") @@ -26,8 +26,10 @@ def tensor2pil(image_tensor): # Helper function to convert PIL Image (RGB) back to ComfyUI tensor (B=1, H, W, C) def pil2tensor(pil_image): - arr = np.array(pil_image).astype(np.float32) / 255.0 # (H, W, 3) - arr = arr[np.newaxis, ...] # (1, H, W, 3) - No transpose needed for BHWC + if pil_image is None: + return None + arr = np.array(pil_image).astype(np.float32) / 255.0 + arr = arr[np.newaxis, ...] return torch.from_numpy(arr) # Main node class @@ -36,12 +38,16 @@ class NanoSeedEdit: def INPUT_TYPES(cls): return { "required": { - "image": ("IMAGE",), # Supports batch/multi-image input "prompt": ("STRING", {"default": "Edit the image according to this prompt.", "multiline": True}), "model": (["nano_banana", "seedream", "flux_kontext_pro", "qwen_edit_plus"],), "fal_key": ("STRING", {"default": "your_fal_key_here"}), }, "optional": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "image3": ("IMAGE",), + "image4": ("IMAGE",), + "image5": ("IMAGE",), "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": 4}), @@ -55,40 +61,43 @@ class NanoSeedEdit: CATEGORY = "image/edit" OUTPUT_NODE = True - def edit_image(self, image, prompt, model, fal_key, width=0, height=0, num_images=1, seed=0): + def edit_image(self, prompt, model, fal_key, image1=None, image2=None, image3=None, image4=None, image5=None, + width=0, height=0, num_images=1, seed=0): if fal_key == "your_fal_key_here": raise ValueError("Please set your fal.ai API key in the node.") - batch_size = image.shape[0] if len(image.shape) == 4 else 1 - # Limit to 5 images max for multi-support models - batch_size = min(batch_size, 5) + # Collect all non-None images + images = [img for img in [image1, image2, image3, image4, image5] if img is not None] + if not images: + raise ValueError("At least one image input must be connected.") + # Convert each to PIL and encode img_data_uris = [] custom_size = (width > 0 and height > 0) - for b in range(batch_size): - single_image = image[b:b + 1] if batch_size > 1 else image - pil_image = tensor2pil(single_image) + for img_tensor in images: + pil_image = tensor2pil(img_tensor) + if pil_image is None: + continue if custom_size and model == "nano_banana": - # Resize each input for NanoBanana pil_image = pil_image.resize((width, height), Image.LANCZOS) - # Encode to base64 data URI buffer = BytesIO() pil_image.save(buffer, format="PNG") img_str = base64.b64encode(buffer.getvalue()).decode() img_data_uris.append(f"data:image/png;base64,{img_str}") + # Enforce single image for Flux Kontext Pro if model == "flux_kontext_pro" and len(img_data_uris) > 1: - raise ValueError("Flux Kontext Pro supports only a single input image. Use one image or select another model.") + raise ValueError("Flux Kontext Pro supports only a single input image. Use only image1.") - # Prepare API endpoint and payload + # Prepare payload based on model if model == "nano_banana": url = "https://fal.run/fal-ai/nano-banana/edit" payload = { "prompt": prompt, - "image_urls": img_data_uris, # Multiple supported + "image_urls": img_data_uris, "num_images": num_images, "output_format": "png", "sync_mode": True, @@ -97,9 +106,9 @@ class NanoSeedEdit: url = "https://fal.run/fal-ai/bytedance/seedream/v4/edit" payload = { "prompt": prompt, - "image_urls": img_data_uris, # Multiple supported (up to 10) - "num_images": 1, # One generation - "max_images": num_images, # Variations + "image_urls": img_data_uris, + "num_images": 1, + "max_images": num_images, "seed": seed, "enable_safety_checker": True, "sync_mode": True, @@ -110,13 +119,13 @@ class NanoSeedEdit: url = "https://fal.run/fal-ai/flux-pro/kontext" payload = { "prompt": prompt, - "image_url": img_data_uris[0], # Single + "image_url": img_data_uris[0], "num_images": num_images, "seed": seed, "output_format": "png", "sync_mode": True, - "guidance_scale": 3.5, # Default - "safety_tolerance": "2", # Default + "guidance_scale": 3.5, + "safety_tolerance": "2", } if custom_size: payload["image_size"] = {"width": width, "height": height} @@ -124,7 +133,7 @@ class NanoSeedEdit: url = "https://fal.run/fal-ai/qwen-image-edit-plus" payload = { "prompt": prompt, - "image_urls": img_data_uris, # Multiple supported + "image_urls": img_data_uris, "num_images": num_images, "seed": seed, "guidance_scale": 4.0, @@ -137,7 +146,7 @@ class NanoSeedEdit: if custom_size: payload["image_size"] = {"width": width, "height": height} - # Make API call + # API call headers = { "Authorization": f"Key {fal_key}", "Content-Type": "application/json", @@ -152,32 +161,29 @@ class NanoSeedEdit: all_edited_tensors = [] - # Process each generated image (limit to num_images) + # Process up to num_images for img_info in api_result["images"][:num_images]: - # Handle data_uri (sync=True) or url img_data = img_info.get("data_uri") or img_info.get("url") if not img_data: continue if img_data.startswith("data:"): - # Parse base64 data URI - header, encoded = img_data.split(",", 1) + _, encoded = img_data.split(",", 1) pil_edited = Image.open(BytesIO(base64.b64decode(encoded))) else: - # Download from URL img_resp = requests.get(img_data) if img_resp.status_code != 200: raise ValueError("Failed to download generated image") pil_edited = Image.open(BytesIO(img_resp.content)) tensor_edited = pil2tensor(pil_edited) - all_edited_tensors.append(tensor_edited) + if tensor_edited is not None: + all_edited_tensors.append(tensor_edited) - # Stack all edited tensors into a batch + # Stack output if all_edited_tensors: batched_output = torch.cat(all_edited_tensors, dim=0) else: - # Fallback: empty batch of 1 dummy image batched_output = torch.zeros((1, 512, 512, 3)) - return (batched_output,) + return (batched_output,) \ No newline at end of file