""" Google Gemini 节点 ComfyUI 自定义节点,用于调用 Gemini Flash 模型进行多模态文本生成 """ import base64 import os import time from typing import Dict, List, Optional, Tuple import torch from ..utils.image_utils import tensor_to_pil, encode_image_to_base64 from ..clients.gemini_flash_client import GeminiFlashClient from ..models_config import get_enabled_flash_models # 支持的视频 MIME 类型映射 VIDEO_MIME_TYPES = { ".mp4": "video/mp4", ".mpeg": "video/mpeg", ".mpg": "video/mpg", ".mov": "video/mov", ".avi": "video/avi", ".flv": "video/x-flv", ".webm": "video/webm", ".wmv": "video/wmv", ".3gp": "video/3gpp", ".3gpp": "video/3gpp" } class GoogleGemini: """ Google Gemini 节点 功能： - 支持多个 Gemini Flash 模型 - 支持图片和视频输入 - 支持系统指令 - 支持不同思考深度（不思考/高） - 输出生成的文本内容 """ # 支持的思考深度选项 THINKING_DEPTHS = ["不思考", "高"] def __init__(self): """初始化节点""" self.client = None @classmethod def INPUT_TYPES(cls): """ 定义输入参数 """ # 从配置获取启用的模型列表 enabled_models = get_enabled_flash_models() default_model = enabled_models[0] if enabled_models else "gemini-3-flash-preview" return { "required": { "模型": (enabled_models, { "default": default_model }), "提示词": ("STRING", { "default": "", "multiline": True }), "思考深度": (cls.THINKING_DEPTHS, { "default": "不思考" }) }, "optional": { "系统指令": ("STRING", { "default": "", "multiline": True }), "图片": ("IMAGE",), "视频": ("VIDEO",) } } # 返回值类型 RETURN_TYPES = ("STRING", "STRING") RETURN_NAMES = ("主要内容", "思考内容") # 执行函数名 FUNCTION = "generate" # 节点分类 CATEGORY = "text/generation" # 允许输出到 UI OUTPUT_NODE = True def _prepare_image_data( self, images: Optional[torch.Tensor] ) -> Optional[List[Dict[str, str]]]: """ 准备图片数据 Args: images: ComfyUI 图片张量 [B, H, W, C] Returns: 图片数据列表，每个元素包含 mime_type 和 data """ if images is None: return None image_data = [] pil_images = tensor_to_pil(images) for img in pil_images: b64_str = encode_image_to_base64(img) image_data.append({ "mime_type": "image/png", "data": b64_str }) return image_data if image_data else None def _prepare_video_data( self, video ) -> Optional[Dict[str, str]]: """ 准备视频数据 ComfyUI VIDEO 类型包含视频文件路径信息。 读取视频文件并转换为 base64。 Args: video: ComfyUI VIDEO 类型数据 Returns: 视频数据字典，包含 mime_type 和 data """ if video is None: return None # VIDEO 类型通常是一个字典，包含 'video' 键指向文件路径 # 或者直接是文件路径字符串 video_path = None if isinstance(video, dict): # 尝试获取视频路径 video_path = video.get("video") or video.get("path") or video.get("file") elif isinstance(video, str): video_path = video elif hasattr(video, "video"): video_path = video.video if not video_path or not os.path.exists(video_path): print(f"Google Gemini: 视频文件不存在或路径无效: {video_path}") return None # 获取文件扩展名和 MIME 类型 _, ext = os.path.splitext(video_path) ext = ext.lower() mime_type = VIDEO_MIME_TYPES.get(ext, "video/mp4") # 检查文件大小（限制 20MB） file_size = os.path.getsize(video_path) if file_size > 20 * 1024 * 1024: raise ValueError( f"视频文件过大 ({file_size / 1024 / 1024:.2f}MB)，" f"请使用不超过 20MB 的视频文件" ) # 读取并编码视频 try: with open(video_path, "rb") as f: video_bytes = f.read() b64_str = base64.b64encode(video_bytes).decode("utf-8") return { "mime_type": mime_type, "data": b64_str } except Exception as e: print(f"Google Gemini: 读取视频文件失败 - {str(e)}") return None def _parse_dual_output(self, raw_response: Dict) -> Tuple[str, str]: """ 解析包含思考内容和主要内容的响应 Args: raw_response: API 原始响应字典 Returns: (主要内容, 思考内容) """ candidates = raw_response.get("candidates", []) if not candidates: return ("", "") parts = candidates[0].get("content", {}).get("parts", []) thought_text = "" main_text = "" for part in parts: if part.get("thought") is True: # 思考部分 thought_text = part.get("text", "") elif "thoughtSignature" in part or "text" in part: # 主要内容 main_text = part.get("text", "") return (main_text, thought_text) def generate( self, 模型: str, 提示词: str, 思考深度: str, 系统指令: Optional[str] = None, 图片: Optional[torch.Tensor] = None, 视频=None ) -> Tuple[str]: """ 生成文本 Args: 模型: 使用的模型名称 提示词: 用户提示词 思考深度: 思考深度选项 系统指令: 系统级指令 图片: 输入图片 视频: 输入视频 Returns: 生成的文本 (STRING,) """ start_time = time.time() try: # 初始化 API 客户端 if self.client is None: try: self.client = GeminiFlashClient() except ValueError as e: raise ValueError(f"初始化失败: {str(e)}") # 准备图片数据 image_data = self._prepare_image_data(图片) if image_data: print(f"Google Gemini: 输入 {len(image_data)} 张图片") # 准备视频数据 video_data = self._prepare_video_data(视频) if video_data: print(f"Google Gemini: 输入视频 ({video_data['mime_type']})") # 构建输入描述 input_desc = [] if 提示词: input_desc.append("文本") if image_data: input_desc.append(f"{len(image_data)}张图片") if video_data: input_desc.append("视频") print(f"Google Gemini: 模型 = {模型}") print(f"Google Gemini: 多模态输入 ({', '.join(input_desc)})") print(f"Google Gemini: 思考深度 = {思考深度}") print(f"Google Gemini: 发送请求...") # 获取端点和构建请求体 endpoint = self.client.get_endpoint(model=模型, thinking_depth=思考深度) request_body = self.client.build_request_body( prompt=提示词, system_instruction=系统指令, image_data=image_data, video_data=video_data ) # 根据是否有视频设置超时 timeout = 300 if video_data else 180 # 调用底层 API 获取原始响应 async def get_raw_response(): return await self.client.request_async( endpoint, request_body, session=None, timeout=timeout ) # 在独立线程中执行异步请求 raw_response = self.client.run_async_in_thread(get_raw_response()) # 计算耗时 elapsed = time.time() - start_time # 解析响应，分离主要内容和思考内容 main_text, thought_text = self._parse_dual_output(raw_response) # 输出信息 print(f"Google Gemini: 生成完成 (耗时: {elapsed:.2f}s)") print(f"Google Gemini: 主要内容长度: {len(main_text)} 字符") print(f"Google Gemini: 思考内容长度: {len(thought_text)} 字符") # 输出预览 if main_text: preview = main_text[:100] + "..." if len(main_text) > 100 else main_text print(f"Google Gemini: 主要内容预览: {preview}") return (main_text, thought_text) except ValueError as e: # 检测是否为授权错误 if str(e) == "未授权！": print("请联系作者授权后方可使用！") else: print(f"Google Gemini: 输入错误 - {str(e)}") raise except RuntimeError as e: print(f"Google Gemini: API 错误 - {str(e)}") raise except Exception as e: print(f"Google Gemini: 未知错误 - {str(e)}") raise