From 78b2951b5f2dfe9fb98d3d4ab57599e0baafcdb5 Mon Sep 17 00:00:00 2001 From: comrender <139063732+comrender@users.noreply.github.com> Date: Wed, 17 Sep 2025 15:07:00 +0200 Subject: [PATCH] Add files via upload --- __init__.py | 2 ++ nano_seed.py | 79 +++++++++++++++++++++++++++++++--------------------- 2 files changed, 50 insertions(+), 31 deletions(-) diff --git a/__init__.py b/__init__.py index 94e7bf5..7d0f56e 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,5 @@ +from .nano_seed import NanoSeedEdit + NODE_CLASS_MAPPINGS = { "NanoSeedEdit": NanoSeedEdit, } diff --git a/nano_seed.py b/nano_seed.py index 3f48ea3..61ac395 100644 --- a/nano_seed.py +++ b/nano_seed.py @@ -4,23 +4,30 @@ import base64 import numpy as np from io import BytesIO from PIL import Image -import os -# Helper function to convert ComfyUI tensor (B, H, W, C) to PIL Image (single or first from batch) +# Helper function to convert ComfyUI tensor (B=1, H, W, C) to PIL Image (RGB) def tensor2pil(image_tensor): - # Assume single image or take first from batch - if image_tensor.shape[0] > 1: - image_tensor = image_tensor[0:1] # Process one at a time if batched here - i = 255. * image_tensor.cpu().numpy() + # image_tensor is always (1, H, W, C) from the loop + i = 255. * image_tensor[0].cpu().numpy() # (H, W, C) image = np.clip(i, 0, 255).astype(np.uint8) - image = np.transpose(image[0], (1, 2, 0)) # HWC - return Image.fromarray(image) + + # Handle channels for PIL (always output RGB) + c = image.shape[-1] + if c == 1: + image = np.repeat(image, 3, axis=-1) # (H, W, 3) + elif c == 3: + pass # Already good + elif c == 4: + image = image[..., :3] # Drop alpha + else: + raise ValueError(f"Unsupported channels: {c}. Expected 1, 3, or 4.") + + return Image.fromarray(image, mode='RGB') -# Helper function to convert PIL Image back to ComfyUI tensor (B=1, H, W, C) +# 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 - arr = arr[np.newaxis, ...] # Add batch dim - arr = np.transpose(arr, (0, 3, 1, 2)) # BHWC + 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 return torch.from_numpy(arr) # Main node class @@ -49,22 +56,26 @@ class NanoSeedEdit: OUTPUT_NODE = True def edit_image(self, image, prompt, model, fal_key, 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 all_edited_tensors = [] for b in range(batch_size): # Extract single image tensor from batch - single_image = image[b:b+1] if batch_size > 1 else image + single_image = image[b:b + 1] if batch_size > 1 else image # Convert to PIL pil_image = tensor2pil(single_image) # Handle custom resolution custom_size = (width > 0 and height > 0) - if custom_size and model == "nano_banana": - # For NanoBanana, resize input image - pil_image = pil_image.resize((width, height), Image.LANCZOS) - # For Seedream, resolution is handled in API payload + if custom_size: + if model == "nano_banana": + # For NanoBanana, resize input image + pil_image = pil_image.resize((width, height), Image.LANCZOS) + # For Seedream, resolution is handled in API payload # Encode to base64 data URI buffer = BytesIO() @@ -87,8 +98,8 @@ class NanoSeedEdit: payload = { "prompt": prompt, "image_urls": [img_data_uri], - "num_images": num_images, - "max_images": 1, # Keep simple; can extend for more variations + "num_images": 1, # One generation + "max_images": num_images, # Variations per generation "seed": seed, "enable_safety_checker": True, "sync_mode": True, @@ -109,17 +120,24 @@ class NanoSeedEdit: if "images" not in api_result or len(api_result["images"]) == 0: raise ValueError("No images returned from API") - # Download and convert each generated image - for img_info in api_result["images"][:num_images]: # Limit to requested num - img_url = img_info.get("url") - if not img_url: + # Process each generated image (limit 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 - img_resp = requests.get(img_url) - if img_resp.status_code != 200: - raise ValueError("Failed to download generated image") + if img_data.startswith("data:"): + # Parse base64 data URI + header, 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)) - pil_edited = Image.open(BytesIO(img_resp.content)) tensor_edited = pil2tensor(pil_edited) all_edited_tensors.append(tensor_edited) @@ -127,8 +145,7 @@ class NanoSeedEdit: if all_edited_tensors: batched_output = torch.cat(all_edited_tensors, dim=0) else: - batched_output = torch.empty((1, 1, 1, 3)) # Fallback empty + # Fallback: empty batch of 1 dummy image + batched_output = torch.zeros((1, 512, 512, 3)) - return (batched_output,) - -# For batch processing in workflows, ComfyUI handles it via the IMAGE type \ No newline at end of file + return (batched_output,) \ No newline at end of file