""" NanoBanana Node V3 - Workflow-Level Parallel Execution with Metadata & Logging Version: 1.0.0 This V3 node enables automatic workflow-graph level parallelism using ComfyUI's async execution engine. Multiple instances will execute concurrently without blocking. Features: - Workflow-level parallel execution via async/await - Real-time usage metadata display on node UI (processing time, token counts) - Per-instance logging to logs/ directory with JSON Lines format - Thoughts output for model reasoning content - Dynamic image inputs (up to 14 images) - Auto aspect ratio detection V3 API Implementation: - Inherits from io.ComfyNode for async support - Uses async def execute() - enables asyncio.Task wrapping by ComfyUI - Uses aiohttp instead of requests - non-blocking HTTP - Uses define_schema() instead of INPUT_TYPES() - Uses ComfyExtension + comfy_entrypoint() for registration - Returns io.NodeOutput instead of tuple - Accesses hidden params via cls.hidden.{name} """ import asyncio import aiohttp import json import time import base64 import io as python_io import os import datetime from pathlib import Path import numpy as np import torch from PIL import Image import hashlib from dotenv import load_dotenv # V3 API imports from comfy_api.latest import io, ui, ComfyExtension # Import PromptServer for metadata display try: from server import PromptServer HAS_SERVER = True except ImportError: HAS_SERVER = False print("[NanoBanana V3] PromptServer not available, metadata display disabled") # Load environment variables load_dotenv() # Constants MAX_IMAGES = 14 SUPPORTED_ASPECT_RATIOS = { "1:1": 1.0, "4:3": 4/3, "3:4": 3/4, "16:9": 16/9, "9:16": 9/16, "3:2": 3/2, "2:3": 2/3, "5:4": 5/4, "4:5": 4/5, "21:9": 21/9, } # ============================================================================ # NodeLogger - Per-Instance Logging # ============================================================================ class NodeLogger: """ Manages per-instance log files for Nano Banana nodes. Each node instance gets its own log file, with workflow runs separated. """ def __init__(self): self.base_path = Path(os.path.dirname(__file__)) / "logs" self.base_path.mkdir(exist_ok=True) def _get_log_filename(self, node_id): """Get the log file path for a specific node instance""" return self.base_path / f"node_{node_id}.log" def log_execution(self, node_id, workflow_run_id, metadata, model, prompt_preview=None, response_preview=None, request_index=None, total_requests=None, is_summary=False): """ Log a node execution with all metadata. """ log_file = self._get_log_filename(node_id) log_entry = { "timestamp": datetime.datetime.now().isoformat(), "workflow_run_id": workflow_run_id, "node_id": node_id, "model": model, "usage_metadata": metadata, "prompt_preview": prompt_preview[:200] if prompt_preview else None, "response_preview": response_preview[:200] if response_preview else None, } # Add parallel execution info if applicable if request_index is not None: log_entry["request_index"] = request_index log_entry["total_requests"] = total_requests if is_summary: log_entry["is_summary"] = True # Append to log file (JSON Lines format for easy parsing) with open(log_file, 'a', encoding='utf-8') as f: f.write(json.dumps(log_entry) + "\n") # Also print detailed console log self._print_console_log(log_entry) def _print_console_log(self, entry): """Print formatted log entry to console""" metadata = entry.get("usage_metadata", {}) request_index = entry.get("request_index") total_requests = entry.get("total_requests") is_summary = entry.get("is_summary", False) # Format header based on type if is_summary: header = "=== Parallel Execution Summary ===" elif request_index is not None: header = f"=== Parallel Request {request_index + 1}/{total_requests} ===" else: header = "=== Execution Log ===" print(f"\n[NanoBanana V3] {header}", flush=True) print(f"[NanoBanana V3] Timestamp: {entry['timestamp']}", flush=True) print(f"[NanoBanana V3] Node ID: {entry['node_id']}", flush=True) print(f"[NanoBanana V3] Workflow Run: {entry['workflow_run_id']}", flush=True) print(f"[NanoBanana V3] Model: {entry['model']}", flush=True) print(f"[NanoBanana V3] --- Usage Metadata ---", flush=True) print(f"[NanoBanana V3] Processing Time: {metadata.get('processingTimeMs', 0):.0f}ms", flush=True) print(f"[NanoBanana V3] Prompt Tokens: {metadata.get('promptTokenCount', 0)}", flush=True) print(f"[NanoBanana V3] Candidates Tokens: {metadata.get('candidatesTokenCount', 0)}", flush=True) print(f"[NanoBanana V3] Thoughts Tokens: {metadata.get('thoughtsTokenCount', 0)}", flush=True) print(f"[NanoBanana V3] Total Tokens: {metadata.get('totalTokenCount', 0)}", flush=True) print(f"[NanoBanana V3] Cached Content Tokens: {metadata.get('cachedContentTokenCount', 0)}", flush=True) if is_summary and total_requests: print(f"[NanoBanana V3] Parallel Requests: {total_requests}", flush=True) print(f"[NanoBanana V3] =========================\n", flush=True) # ============================================================================ # Utility Functions # ============================================================================ def calculate_aspect_ratio(image): """Calculate the closest standard aspect ratio from an image tensor.""" import math width, height = None, None if isinstance(image, torch.Tensor): if image.ndim == 4: height, width = int(image.shape[1]), int(image.shape[2]) elif image.ndim == 3: height, width = int(image.shape[0]), int(image.shape[1]) elif isinstance(image, np.ndarray): if image.ndim == 4: height, width = int(image.shape[1]), int(image.shape[2]) elif image.ndim == 3: height, width = int(image.shape[0]), int(image.shape[1]) if width is None or height is None or height <= 0 or width <= 0: return "1:1" actual_ratio = width / height log_actual = math.log(actual_ratio) closest_name = "1:1" min_log_diff = float('inf') for name, ratio_value in SUPPORTED_ASPECT_RATIOS.items(): log_diff = abs(log_actual - math.log(ratio_value)) if log_diff < min_log_diff: min_log_diff = log_diff closest_name = name return closest_name def image_to_base64(image): """Convert ComfyUI IMAGE tensor to base64 PNG string.""" if image.ndim == 4: image = image.squeeze(0) image_np = image.cpu().numpy() if image_np.dtype != np.uint8: image_np = np.clip(image_np, 0, 1) image_np = (image_np * 255).astype(np.uint8) pil_image = Image.fromarray(image_np, 'RGB') buffered = python_io.BytesIO() pil_image.save(buffered, format="PNG") return base64.b64encode(buffered.getvalue()).decode('utf-8') def base64_to_image(base64_str): """Convert base64 image string to ComfyUI tensor [1, H, W, 3].""" try: img_data = base64.b64decode(base64_str) img = Image.open(python_io.BytesIO(img_data)).convert("RGB") img_array = np.array(img).astype(np.float32) / 255.0 return torch.from_numpy(img_array).unsqueeze(0) except Exception as e: print(f"[NanoBanana V3] Error in base64_to_image: {e}") return torch.zeros((1, 64, 64, 3), dtype=torch.float32) # ============================================================================ # V3 Node Definition # ============================================================================ class NanoBananaNodeV3(io.ComfyNode): """ V3 Node for Vertex AI's generateContent API with workflow-level parallelism. Features: - Async execution enables multiple nodes to run concurrently in workflow - Usage metadata display on node UI - Per-instance logging to logs/ directory - Thoughts output for model reasoning - Proper timing for async operations """ # Class-level cache for models _models_cache = None _last_fetch_time = 0 _cache_duration = 3600 @classmethod def _get_models(cls): """Fetch available models with caching.""" current_time = time.time() if cls._models_cache is None or (current_time - cls._last_fetch_time > cls._cache_duration): models_env = os.getenv("VERTEX_AI_MODELS", "gemini-3-pro-image-preview,gemini-2.5-flash-image") cls._models_cache = sorted([m.strip() for m in models_env.split(",") if m.strip()]) cls._last_fetch_time = current_time return cls._models_cache or ["gemini-3-pro-image-preview"] @classmethod def define_schema(cls): """ V3 schema definition - replaces INPUT_TYPES, RETURN_TYPES, CATEGORY, etc. """ return io.Schema( node_id="NanoBananaNodeV3", display_name="Nano Banana (Pro) V3", description="Vertex AI image generation with workflow-level parallel execution, metadata display, and logging", category="LLM", inputs=[ io.String.Input( "system_prompt", multiline=True, default="You are a master artist and expert at digital art." ), io.String.Input( "user_message_box", multiline=True, default="Combine all the input images with a background of your envisioning in the year 2050." ), io.Combo.Input( "model", options=cls._get_models(), default="gemini-3-pro-image-preview" ), io.Boolean.Input( "image_generation", default=True ), io.Combo.Input( "resolution", options=["1K", "2K", "4K"], default="4K" ), io.Combo.Input( "aspect_ratio", options=["None", "Auto", "1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "5:4", "4:5", "21:9"], default="Auto" ), io.Float.Input( "temperature", default=1.0, min=0.0, max=1.0, step=0.1 ), io.Int.Input( "timeout", default=300, min=30, max=600, step=10 ), # Optional image inputs (3 fixed, dynamic inputs added via JS) io.Image.Input("image_1", optional=True), io.Image.Input("image_2", optional=True), io.Image.Input("image_3", optional=True), ], outputs=[ io.String.Output("text_output", display_name="Text Output"), io.Image.Output("image_output", display_name="Image"), io.String.Output("stats", display_name="Stats"), io.String.Output("thoughts", display_name="Thoughts"), ], hidden=[ io.Hidden.unique_id, # Required for logging and metadata display ], is_output_node=False, ) @classmethod async def execute( cls, system_prompt, user_message_box, model, image_generation, resolution, aspect_ratio, temperature, timeout, image_1=None, image_2=None, image_3=None, **kwargs, # Capture dynamic image inputs (image_4, image_5, etc.) AND hidden params ): """ Async execute method - enables workflow-level parallel execution. ComfyUI's execution engine wraps this in asyncio.Task, allowing multiple independent nodes to execute concurrently. """ placeholder_image = torch.zeros((1, 1, 1, 3), dtype=torch.float32) # V3 nodes access hidden params via cls.hidden.{name} unique_id = cls.hidden.unique_id if hasattr(cls, 'hidden') and hasattr(cls.hidden, 'unique_id') else None # Generate workflow run ID for logging workflow_run_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f") # Get API key api_key = os.getenv("VERTEX_AI_API_KEY", "") if not api_key: return io.NodeOutput( "Error: VERTEX_AI_API_KEY not found", placeholder_image, "Stats N/A", "" ) # Build endpoint URL endpoint_base = "https://aiplatform.googleapis.com/v1" use_simple_endpoint = os.getenv("VERTEX_AI_USE_SIMPLE_ENDPOINT", "false").lower() == "true" if use_simple_endpoint: endpoint = os.getenv("VERTEX_AI_ENDPOINT") if not endpoint: return io.NodeOutput( "Error: VERTEX_AI_ENDPOINT must be set", placeholder_image, "Stats N/A", "" ) url = f"{endpoint}/google/{model}:generateContent" else: project = os.getenv("VERTEX_AI_PROJECT", "project") location = os.getenv("VERTEX_AI_LOCATION", "location") url = f"{endpoint_base}/projects/{project}/locations/{location}/publishers/*/models/{model}:generateContent" headers = { "api-key": api_key, "Content-Type": "application/json", } # Validate temperature (max 1.0 to match V1) validated_temp = max(0.0, min(1.0, float(temperature))) # Process input images (3 fixed + dynamic from kwargs) image_data = [] fixed_images = [img for img in [image_1, image_2, image_3] if img is not None] dynamic_images = [kwargs[k] for k in sorted(kwargs.keys()) if k.startswith('image_') and kwargs[k] is not None] input_images = fixed_images + dynamic_images # Calculate aspect ratio if Auto final_aspect_ratio = aspect_ratio if aspect_ratio == "Auto" and input_images: final_aspect_ratio = calculate_aspect_ratio(input_images[0]) # Send aspect ratio to frontend for display if HAS_SERVER and unique_id: PromptServer.instance.send_sync("nanobanana.aspect_ratio", { "node_id": str(unique_id), "aspect_ratio": final_aspect_ratio, "was_auto": True, "node_type": "NanoBananaNodeV3" }) elif aspect_ratio == "Auto": final_aspect_ratio = "None" # Convert images to base64 for i, img in enumerate(input_images): try: img_str = image_to_base64(img) image_data.append({ "inlineData": { "mimeType": "image/png", "data": img_str } }) except Exception as e: return io.NodeOutput( f"Error processing image_{i+1}: {e}", placeholder_image, "Stats N/A", "" ) # Build request payload user_parts = [{"text": user_message_box}] user_parts.extend(image_data) data = { "contents": [{"role": "user", "parts": user_parts}], "systemInstruction": {"parts": [{"text": system_prompt}]}, "generationConfig": {"temperature": validated_temp} } # Add image generation config if image_generation: data["generationConfig"]["responseModalities"] = ["TEXT", "IMAGE"] data["generationConfig"]["candidateCount"] = 1 image_config = {"image_size": resolution} if final_aspect_ratio != "None": image_config["aspect_ratio"] = final_aspect_ratio data["generationConfig"]["image_config"] = image_config # ================================================================ # ASYNC HTTP REQUEST - Non-blocking! # Using aiohttp allows the event loop to handle other tasks # while waiting for the API response. # ================================================================ try: start_time = time.time() # Create timeout for aiohttp client_timeout = aiohttp.ClientTimeout(total=timeout) async with aiohttp.ClientSession(timeout=client_timeout) as session: async with session.post(url, headers=headers, json=data) as response: response.raise_for_status() result = await response.json() end_time = time.time() elapsed_time = end_time - start_time processing_time_ms = elapsed_time * 1000 # Process response if not result.get("candidates") or not result["candidates"][0].get("content"): return io.NodeOutput( "Error: Invalid response format from API", placeholder_image, "Stats N/A", "" ) candidate = result["candidates"][0] parts = candidate.get("content", {}).get("parts", []) text_output = "" image_tensor = placeholder_image thoughts_output = "" # Extract text, thoughts, and images from parts for part in parts: if part.get("thought") == True and "text" in part: # Model's internal reasoning/thinking thought_text = part["text"] if isinstance(thought_text, str): thoughts_output += thought_text elif "text" in part and isinstance(part["text"], str): # Regular text output text_output += part["text"] elif "inlineData" in part: inline_data = part["inlineData"] if inline_data.get("mimeType", "").startswith("image"): image_tensor = base64_to_image(inline_data.get("data", "")) # Extract full usage metadata usage_metadata = result.get("usageMetadata", {}) prompt_tokens = usage_metadata.get("promptTokenCount", 0) completion_tokens = usage_metadata.get("candidatesTokenCount", 0) tps = completion_tokens / elapsed_time if elapsed_time > 0 else 0 full_metadata = { "processingTimeMs": processing_time_ms, "promptTokenCount": usage_metadata.get("promptTokenCount", 0), "candidatesTokenCount": usage_metadata.get("candidatesTokenCount", 0), "totalTokenCount": usage_metadata.get("totalTokenCount", 0), "thoughtsTokenCount": usage_metadata.get("thoughtsTokenCount", 0), "cachedContentTokenCount": usage_metadata.get("cachedContentTokenCount", 0), } # Send usage metadata to frontend for display if HAS_SERVER and unique_id: PromptServer.instance.send_sync("nanobanana.usage_metadata", { "node_id": str(unique_id), "metadata": full_metadata, "model": model, "node_type": "NanoBananaNodeV3" }) # Log to per-instance log file if unique_id: logger = NodeLogger() logger.log_execution( node_id=str(unique_id), workflow_run_id=workflow_run_id, metadata=full_metadata, model=model, prompt_preview=user_message_box if user_message_box else None, response_preview=text_output if text_output else None ) stats_text = f"TPS: {tps:.2f}, Time: {elapsed_time:.1f}s, Tokens: {prompt_tokens + completion_tokens}" return io.NodeOutput( text_output, image_tensor, stats_text, thoughts_output ) except aiohttp.ClientError as e: error_msg = f"API Request Error: {str(e)}" print(f"[NanoBanana V3] {error_msg}") return io.NodeOutput(error_msg, placeholder_image, "Stats N/A", "") except asyncio.TimeoutError: error_msg = f"Request timed out after {timeout}s" print(f"[NanoBanana V3] {error_msg}") return io.NodeOutput(error_msg, placeholder_image, "Stats N/A", "") except Exception as e: import traceback print(f"[NanoBanana V3] Exception: {e}") print(f"[NanoBanana V3] Traceback: {traceback.format_exc()}") return io.NodeOutput(f"Node Error: {str(e)}", placeholder_image, "Stats N/A", "") @classmethod def fingerprint_inputs(cls, **kwargs): """ V3 equivalent of IS_CHANGED - controls caching behavior. Returns a fingerprint that changes when inputs change. """ image_hashes = [] # Get all image keys (fixed + dynamic) image_keys = sorted([k for k in kwargs.keys() if k.startswith('image_')]) for key in image_keys: img = kwargs.get(key) if img is not None and isinstance(img, torch.Tensor): try: hasher = hashlib.sha256() hasher.update(img.cpu().numpy().tobytes()) image_hashes.append(hasher.hexdigest()) except Exception: image_hashes.append(f"{key}_error") return ( kwargs.get("system_prompt"), kwargs.get("user_message_box"), kwargs.get("model"), kwargs.get("image_generation"), kwargs.get("resolution"), kwargs.get("aspect_ratio"), float(kwargs.get("temperature", 1.0)), tuple(image_hashes) ) # ============================================================================ # V3 Extension Registration # ============================================================================ class NanoBananaExtension(ComfyExtension): """ V3 Extension class - replaces NODE_CLASS_MAPPINGS / NODE_DISPLAY_NAME_MAPPINGS """ async def get_node_list(self): """Return list of nodes provided by this extension.""" return [ NanoBananaNodeV3, ] async def comfy_entrypoint(): """ V3 entry point - ComfyUI calls this to load the extension. Replaces the old NODE_CLASS_MAPPINGS dictionary pattern. """ return NanoBananaExtension()