""" PD: Gemini Image Generation Node (ComfyUI Auth Token) 基于 gemini_image_api 的 Auth Token 认证方式 支持文生图和图生图 """ import torch import requests import time import io import base64 import numpy as np from PIL import Image # 默认系统提示词（来自 ComfyUI 官方 API 节点） GEMINI_IMAGE_SYS_PROMPT = ( "You are an expert image-generation engine. You must ALWAYS produce an image.\n" "Interpret all user input—regardless of " "format, intent, or abstraction—as literal visual directives for image composition.\n" "If a prompt is conversational or lacks specific visual details, " "you must creatively invent a concrete visual scenario that depicts the concept.\n" "Prioritize generating the visual representation above any text, formatting, or conversational requests." ) def tensor_to_base64(tensor): """将 Tensor 转换为 base64 字符串""" if tensor.ndim == 4: tensor = tensor[0] array = np.clip(255. * tensor.cpu().numpy(), 0, 255).astype(np.uint8) image = Image.fromarray(array) buffered = io.BytesIO() image.save(buffered, format="PNG") return base64.b64encode(buffered.getvalue()).decode("utf-8") def bytes_to_tensor(data): """将图片字节数据转换为 Tensor""" try: img = Image.open(io.BytesIO(data)).convert("RGB") image_np = np.array(img).astype(np.float32) / 255.0 return torch.from_numpy(image_np).unsqueeze(0) except Exception as e: print(f"[GEMINI_AUTHTOKEN] Image decode error: {e}") return torch.zeros((1, 512, 512, 3)) class PDGeminiImageGenAuthToken: """ PD: Gemini Image Generation (Auth Token) 使用 ComfyUI Auth Token 调用 Gemini Image API """ @classmethod def INPUT_TYPES(cls): return { "required": { "auth_token": ("STRING", { "default": "", "multiline": False, "placeholder": "Paste your Auth Token here" }), "prompt": ("STRING", { "multiline": True, "default": "A futuristic city with flying cars", "tooltip": "Describe what you want to generate" }), "model": ([ "gemini-2.5-flash-image", "gemini-2.5-flash-image-preview" ], { "default": "gemini-2.5-flash-image" }), "aspect_ratio": ([ "auto", "1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21" ], { "default": "auto" }), "resolution": (["auto", "1K", "2K", "4K"], { "default": "1K" }), "seed": ("INT", { "default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF }), }, "optional": { "images": ("IMAGE", { "tooltip": "Reference image(s) for image-to-image generation" }), "files": ("GEMINI_INPUT_FILES", ), "system_prompt": ("STRING", { "multiline": True, "default": GEMINI_IMAGE_SYS_PROMPT, "tooltip": "Optional system instructions" }), }, "hidden": { "unique_id": "UNIQUE_ID", } } RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("image", "info") FUNCTION = "generate_image" CATEGORY = "PD_Tools/Image_Generation" def generate_image(self, auth_token, prompt, model, aspect_ratio, resolution, seed, unique_id, images=None, files=None, system_prompt=""): """生成图片""" print(f"[GEMINI_AUTHTOKEN] ========== 开始执行 ==========") print(f"[GEMINI_AUTHTOKEN] Prompt: {prompt[:50]}...") print(f"[GEMINI_AUTHTOKEN] Model: {model}") # 验证 Token if not auth_token or not auth_token.strip(): error_msg = "Error: Auth Token is required" print(f"[GEMINI_AUTHTOKEN] {error_msg}") return (torch.zeros((1, 512, 512, 3)), error_msg) # 清理 Token auth_token = auth_token.strip() if auth_token.lower().startswith("bearer "): auth_token = auth_token[7:].strip() # 构建 API URL (不需要添加 :generateContent 后缀，API 会自动处理) base_url = f"https://api.comfy.org/proxy/vertexai/gemini/{model}" # 构建请求内容 parts = [] # 添加系统提示词（如果有） if system_prompt and system_prompt.strip(): parts.append({"text": system_prompt.strip()}) # 添加用户提示词 parts.append({"text": prompt}) # 添加参考图片（如果有） if images is not None: for i in range(images.shape[0]): b64_str = tensor_to_base64(images[i]) parts.append({ "inlineData": { "mimeType": "image/png", "data": b64_str } }) print(f"[GEMINI_AUTHTOKEN] Added reference image {i+1}") # Files 输入处理 if files: for file_part in files: if hasattr(file_part, "model_dump"): parts.append(file_part.model_dump(exclude_none=True)) else: parts.append(file_part) print(f"[GEMINI_AUTHTOKEN] Added file part") # 生成配置 generation_config = { "responseModalities": ["IMAGE"] } # 图片配置 image_config = {} if aspect_ratio != "auto": image_config["aspectRatio"] = aspect_ratio if resolution != "auto": image_config["resolution"] = resolution if image_config: generation_config["imageConfig"] = image_config # 完整请求体 payload = { "contents": [{ "role": "user", "parts": parts }], "generationConfig": generation_config } # 请求头 headers = { "Authorization": f"Bearer {auth_token}", "Content-Type": "application/json", "User-Agent": "ComfyUI-PD-Node/2.0" } start_time = time.time() try: print(f"[GEMINI_AUTHTOKEN] Sending request to {model}") print(f"[GEMINI_AUTHTOKEN] Aspect Ratio: {aspect_ratio}, Resolution: {resolution}") response = requests.post(base_url, headers=headers, json=payload, timeout=120) # 检查响应 if response.status_code != 200: error_msg = f"API Error {response.status_code}: {response.text[:200]}" print(f"[GEMINI_AUTHTOKEN] {error_msg}") return (torch.zeros((1, 512, 512, 3)), error_msg) result = response.json() # 提取图片 image_tensors = [] candidates = result.get("candidates", []) if not candidates: error_msg = "Error: No candidates in response" print(f"[GEMINI_AUTHTOKEN] {error_msg}") return (torch.zeros((1, 512, 512, 3)), error_msg) for cand in candidates: content = cand.get("content", {}) parts = content.get("parts", []) for part in parts: # 支持两种字段名格式 inline_data = part.get("inline_data") or part.get("inlineData") if inline_data: mime = inline_data.get("mime_type") or inline_data.get("mimeType") data_b64 = inline_data.get("data") if mime in ["image/png", "image/jpeg"] and data_b64: raw_data = base64.b64decode(data_b64) image_tensors.append(bytes_to_tensor(raw_data)) print(f"[GEMINI_AUTHTOKEN] Extracted image ({mime})") if not image_tensors: error_msg = f"Error: No images found in response" print(f"[GEMINI_AUTHTOKEN] {error_msg}") print(f"[GEMINI_AUTHTOKEN] Response: {str(result)[:500]}") return (torch.zeros((1, 512, 512, 3)), error_msg) final_image = torch.cat(image_tensors, dim=0) # 计算耗时和成本 end_time = time.time() duration = end_time - start_time # 估算成本（根据模型） cost_per_img = 0.04 if "pro" in model.lower(): cost_per_img = 0.08 total_cost = cost_per_img * len(image_tensors) info_str = ( f"Model: {model}\n" f"Status: Success\n" f"Time: {duration:.2f}s\n" f"Images: {len(image_tensors)}\n" f"Aspect Ratio: {aspect_ratio}\n" f"Resolution: {resolution}\n" f"Est. Cost: ${total_cost:.4f}" ) print(f"[GEMINI_AUTHTOKEN] Success! Generated {len(image_tensors)} images in {duration:.2f}s") return (final_image, info_str) except requests.exceptions.Timeout: error_msg = "Error: Request timeout (120s)" print(f"[GEMINI_AUTHTOKEN] {error_msg}") return (torch.zeros((1, 512, 512, 3)), error_msg) except requests.exceptions.ConnectionError as e: error_msg = f"Error: Connection failed - {str(e)}" print(f"[GEMINI_AUTHTOKEN] {error_msg}") return (torch.zeros((1, 512, 512, 3)), error_msg) except Exception as e: error_msg = f"Error: {str(e)}" print(f"[GEMINI_AUTHTOKEN] {error_msg}") import traceback traceback.print_exc() return (torch.zeros((1, 512, 512, 3)), error_msg) # 节点注册 NODE_CLASS_MAPPINGS = { "PDGeminiImageGenAuthToken": PDGeminiImageGenAuthToken } NODE_DISPLAY_NAME_MAPPINGS = { "PDGeminiImageGenAuthToken": "PD: Gemini Image (ComfyUI AuthToken)" }