792 lines
30 KiB
Python
792 lines
30 KiB
Python
import os
|
|
import json
|
|
import base64
|
|
import requests
|
|
from io import BytesIO
|
|
from typing import Optional, List, Dict, Any
|
|
from PIL import Image
|
|
import torch
|
|
import numpy as np
|
|
from typing import List, Union
|
|
|
|
|
|
def pil2tensor(image: Union[Image.Image, List[Image.Image]]) -> torch.Tensor:
|
|
"""
|
|
Convert PIL image(s) to tensor, matching ComfyUI's implementation.
|
|
|
|
Args:
|
|
image: Single PIL Image or list of PIL Images
|
|
|
|
Returns:
|
|
torch.Tensor: Image tensor with values normalized to [0, 1]
|
|
"""
|
|
if isinstance(image, list):
|
|
if len(image) == 0:
|
|
return torch.empty(0)
|
|
return torch.cat([pil2tensor(img) for img in image], dim=0)
|
|
|
|
if image.mode == 'RGBA':
|
|
image = image.convert('RGB')
|
|
elif image.mode != 'RGB':
|
|
image = image.convert('RGB')
|
|
|
|
img_array = np.array(image).astype(np.float32) / 255.0
|
|
|
|
return torch.from_numpy(img_array)[None,]
|
|
|
|
|
|
def tensor2pil(image: torch.Tensor) -> List[Image.Image]:
|
|
"""
|
|
Convert tensor to PIL image(s), matching ComfyUI's implementation.
|
|
|
|
Args:
|
|
image: Tensor with shape [B, H, W, 3] or [H, W, 3], values in range [0, 1]
|
|
|
|
Returns:
|
|
List[Image.Image]: List of PIL Images
|
|
"""
|
|
batch_count = image.size(0) if len(image.shape) > 3 else 1
|
|
if batch_count > 1:
|
|
out = []
|
|
for i in range(batch_count):
|
|
out.extend(tensor2pil(image[i]))
|
|
return out
|
|
|
|
numpy_image = np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
|
|
|
return [Image.fromarray(numpy_image)]
|
|
|
|
|
|
# ==================== 配置管理 ====================
|
|
CONFIG_FILE = os.path.join(os.path.dirname(__file__), "gemini_config.json")
|
|
DEFAULT_BASE_URL = "https://generativelanguage.googleapis.com/v1beta"
|
|
|
|
|
|
def normalize_base_url(base_url: str) -> str:
|
|
"""Normalize base_url for consistent storage and request building."""
|
|
if not isinstance(base_url, str):
|
|
return ""
|
|
return base_url.strip().rstrip("/")
|
|
|
|
|
|
def describe_base_url(base_url: str) -> str:
|
|
"""Return a generic label for status output without exposing the raw URL."""
|
|
normalized_base_url = normalize_base_url(base_url)
|
|
default_base_url = normalize_base_url(DEFAULT_BASE_URL)
|
|
if not normalized_base_url or normalized_base_url == default_base_url:
|
|
return "default"
|
|
return "custom"
|
|
|
|
|
|
def get_config() -> Dict[str, Any]:
|
|
"""读取配置文件"""
|
|
default_config = {
|
|
"api_key": "",
|
|
"base_url": DEFAULT_BASE_URL,
|
|
"key_base_url_map": {},
|
|
}
|
|
|
|
if os.path.exists(CONFIG_FILE):
|
|
try:
|
|
with open(CONFIG_FILE, 'r') as f:
|
|
loaded_config = json.load(f)
|
|
if isinstance(loaded_config, dict):
|
|
config = {**default_config, **loaded_config}
|
|
if not isinstance(config.get("key_base_url_map"), dict):
|
|
config["key_base_url_map"] = {}
|
|
config["base_url"] = normalize_base_url(config.get("base_url", "")) or DEFAULT_BASE_URL
|
|
return config
|
|
except Exception as e:
|
|
print(f"[GoogleAPI] 读取配置失败: {e}")
|
|
return default_config.copy()
|
|
|
|
|
|
def save_config(config: Dict[str, Any]):
|
|
"""保存配置文件"""
|
|
try:
|
|
config_to_save = dict(config)
|
|
config_to_save["base_url"] = normalize_base_url(config_to_save.get("base_url", "")) or DEFAULT_BASE_URL
|
|
|
|
key_base_url_map = config_to_save.get("key_base_url_map", {})
|
|
if not isinstance(key_base_url_map, dict):
|
|
key_base_url_map = {}
|
|
|
|
config_to_save["key_base_url_map"] = {
|
|
str(key).strip(): normalize_base_url(value)
|
|
for key, value in key_base_url_map.items()
|
|
if str(key).strip() and normalize_base_url(value)
|
|
}
|
|
|
|
os.makedirs(os.path.dirname(CONFIG_FILE), exist_ok=True)
|
|
with open(CONFIG_FILE, 'w') as f:
|
|
json.dump(config_to_save, f, indent=2)
|
|
except Exception as e:
|
|
print(f"[GoogleAPI] 保存配置失败: {e}")
|
|
|
|
|
|
# ==================== Google API 客户端 ====================
|
|
class GoogleGeminiClient:
|
|
"""Google Gemini API 客户端"""
|
|
|
|
def __init__(self, api_key: str = "", base_url: str = ""):
|
|
config = get_config()
|
|
self.api_key = api_key.strip() or config.get('api_key', '').strip()
|
|
# self.base_url = "https://generativelanguage.googleapis.com/v1beta"
|
|
self.base_url = self._resolve_base_url(base_url, config)
|
|
self.timeout = 30000 # 请求超时时间，单位秒
|
|
|
|
def _resolve_base_url(self, base_url: str = "", config: Optional[Dict[str, Any]] = None) -> str:
|
|
normalized_runtime_url = normalize_base_url(base_url)
|
|
if normalized_runtime_url:
|
|
return normalized_runtime_url
|
|
|
|
config = config or get_config()
|
|
if self.api_key:
|
|
key_base_url_map = config.get("key_base_url_map", {})
|
|
if isinstance(key_base_url_map, dict):
|
|
mapped_url = normalize_base_url(key_base_url_map.get(self.api_key, ""))
|
|
if mapped_url:
|
|
return mapped_url
|
|
|
|
return normalize_base_url(config.get("base_url", "")) or DEFAULT_BASE_URL
|
|
|
|
def _save_credentials(self):
|
|
if not self.api_key:
|
|
return
|
|
|
|
config = get_config()
|
|
config['api_key'] = self.api_key
|
|
config['base_url'] = self.base_url
|
|
|
|
key_base_url_map = config.get("key_base_url_map", {})
|
|
if not isinstance(key_base_url_map, dict):
|
|
key_base_url_map = {}
|
|
key_base_url_map[self.api_key] = self.base_url
|
|
config["key_base_url_map"] = key_base_url_map
|
|
save_config(config)
|
|
|
|
def set_api_key(self, api_key: str, base_url: str = ""):
|
|
"""设置API密钥"""
|
|
self.api_key = api_key.strip()
|
|
if not self.api_key:
|
|
return
|
|
|
|
self.base_url = self._resolve_base_url(base_url)
|
|
self._save_credentials()
|
|
|
|
def set_base_url(self, base_url: str):
|
|
"""设置base_url"""
|
|
normalized_base_url = normalize_base_url(base_url)
|
|
if not normalized_base_url:
|
|
return
|
|
|
|
self.base_url = normalized_base_url
|
|
|
|
if self.api_key:
|
|
self._save_credentials()
|
|
else:
|
|
config = get_config()
|
|
config['base_url'] = self.base_url
|
|
save_config(config)
|
|
|
|
def apply_runtime_config(self, api_key: str = "", base_url: str = ""):
|
|
"""Apply runtime api_key/base_url and persist key-url pair when key is provided."""
|
|
normalized_api_key = api_key.strip()
|
|
normalized_base_url = normalize_base_url(base_url)
|
|
|
|
if normalized_api_key:
|
|
self.set_api_key(normalized_api_key, normalized_base_url)
|
|
elif normalized_base_url:
|
|
self.set_base_url(normalized_base_url)
|
|
|
|
def _get_headers(self) -> Dict[str, str]:
|
|
"""获取请求头"""
|
|
return {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
}
|
|
|
|
def _build_url(self, model: str, endpoint: str) -> str:
|
|
"""构建API URL"""
|
|
return f"{self.base_url}/models/{model}:{endpoint}?key={self.api_key}"
|
|
|
|
def generate_content(self, model: str, contents: List[Dict],
|
|
generation_config: Optional[Dict] = None,
|
|
system_instruction: Optional[str] = None) -> Dict[str, Any]:
|
|
"""
|
|
调用 generateContent API
|
|
|
|
Args:
|
|
model: 模型名称 (e.g., "gemini-2.0-flash")
|
|
contents: 内容列表
|
|
generation_config: 生成配置
|
|
system_instruction: 系统指令
|
|
|
|
Returns:
|
|
API响应
|
|
"""
|
|
if not self.api_key:
|
|
raise ValueError("API key not set")
|
|
|
|
payload = {
|
|
"contents": contents
|
|
}
|
|
|
|
if generation_config:
|
|
payload["generationConfig"] = generation_config
|
|
|
|
if system_instruction:
|
|
payload["systemInstruction"] = {
|
|
"parts": [{"text": system_instruction}]
|
|
}
|
|
|
|
url = self._build_url(model, "generateContent")
|
|
|
|
try:
|
|
response = requests.post(
|
|
url,
|
|
headers=self._get_headers(),
|
|
json=payload,
|
|
timeout=self.timeout
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except requests.exceptions.RequestException as e:
|
|
raise Exception(f"API request failed: {str(e)}")
|
|
|
|
|
|
# ==================== ComfyUI 节点 ====================
|
|
|
|
class GoogleGeminiImageAnalysis:
|
|
"""Google Gemini 图像分析节点"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True}),
|
|
"model": ("STRING", {"default": "gemini-3.1-pro-preview"}),
|
|
},
|
|
"optional": {
|
|
"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
"image3": ("IMAGE",),
|
|
"image4": ("IMAGE",),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
"base_url": ("STRING", {"default": ""}),
|
|
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
|
|
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"max_tokens": ("INT", {"default": 8192, "min": 1, "max": 65536}),
|
|
"system_instruction": ("STRING", {"default": "", "multiline": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("analysis",)
|
|
FUNCTION = "analyze"
|
|
CATEGORY = "Google/Gemini"
|
|
|
|
def __init__(self):
|
|
self.client = GoogleGeminiClient()
|
|
|
|
def _tensor_to_base64(self, image_tensor: torch.Tensor) -> str:
|
|
"""将张量转换为base64"""
|
|
pil_images = tensor2pil(image_tensor)
|
|
pil_image = pil_images[0]
|
|
|
|
buffered = BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
return base64.b64encode(buffered.getvalue()).decode('utf-8')
|
|
|
|
def analyze(self, prompt: str, model: str, temperature: float, top_p: float,
|
|
max_tokens: int, api_key: str = "", base_url: str = "", system_instruction: str = "",
|
|
image1=None, image2=None, image3=None, image4=None) -> tuple:
|
|
"""分析多个图像"""
|
|
try:
|
|
if api_key.strip() or base_url.strip():
|
|
self.client.apply_runtime_config(api_key=api_key, base_url=base_url)
|
|
|
|
if not self.client.api_key:
|
|
return ("Error: API key not provided",)
|
|
|
|
# 构建parts
|
|
parts = []
|
|
|
|
# 添加所有图像
|
|
for img in [image1, image2, image3, image4]:
|
|
if img is not None:
|
|
image_base64 = self._tensor_to_base64(img)
|
|
parts.append({
|
|
"inlineData": {
|
|
"mimeType": "image/png",
|
|
"data": image_base64
|
|
}
|
|
})
|
|
|
|
# 添加提示词
|
|
parts.append({"text": prompt})
|
|
|
|
# 构建内容
|
|
contents = [{
|
|
"role": "user",
|
|
"parts": parts
|
|
}]
|
|
|
|
# 生成配置
|
|
generation_config = {
|
|
"temperature": temperature,
|
|
"topP": top_p,
|
|
"maxOutputTokens": max_tokens,
|
|
}
|
|
|
|
# 调用API
|
|
response = self.client.generate_content(
|
|
model=model,
|
|
contents=contents,
|
|
generation_config=generation_config,
|
|
system_instruction=system_instruction if system_instruction else None
|
|
)
|
|
|
|
# 提取文本
|
|
if "candidates" in response and response["candidates"]:
|
|
candidate = response["candidates"][0]
|
|
response_parts = candidate.get("content", {}).get("parts") or []
|
|
for part in response_parts:
|
|
# 跳过thought部分
|
|
if part.get("thought", False):
|
|
print("[GoogleAPI] 跳过thought部分")
|
|
continue
|
|
|
|
if "text" in part:
|
|
text = part["text"]
|
|
return (text,)
|
|
else:
|
|
return ("Error: No response from API",)
|
|
|
|
except Exception as e:
|
|
return (f"Error: {str(e)}",)
|
|
|
|
|
|
class GoogleGeminiImageGeneration:
|
|
"""Google Gemini 图像生成节点"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True}),
|
|
"model": (["gemini-2.5-flash-image", "gemini-3-pro-image-preview", "gemini-3.1-flash-image-preview"], {"default": "gemini-2.5-flash-image"}),
|
|
},
|
|
"optional": {
|
|
"api_key": ("STRING", {"default": ""}),
|
|
"base_url": ("STRING", {"default": ""}),
|
|
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
|
|
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"max_tokens": ("INT", {"default": 8192, "min": 1, "max": 65536}),
|
|
"system_instruction": ("STRING", {"default": "", "multiline": True}),
|
|
"aspect_ratio": (["auto", "16:9", "4:3", "4:5", "3:2", "1:1", "2:3", "3:4", "5:4", "9:16", "21:9"], {"default": "auto"}),
|
|
"image_size": (["auto", "1K", "2K", "4K"], {"default": "auto"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 2147483647}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "STRING")
|
|
RETURN_NAMES = ("image", "description", "response_info")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Google/Gemini"
|
|
|
|
def __init__(self):
|
|
self.client = GoogleGeminiClient()
|
|
|
|
def _base64_to_tensor(self, base64_str: str) -> torch.Tensor:
|
|
"""将base64转换为张量"""
|
|
try:
|
|
image_data = base64.b64decode(base64_str)
|
|
image = Image.open(BytesIO(image_data)).convert('RGB')
|
|
return pil2tensor(image)
|
|
except Exception as e:
|
|
print(f"[GoogleAPI] Base64转换失败: {e}")
|
|
return None
|
|
|
|
def generate(self, prompt: str, model: str, temperature: float, top_p: float,
|
|
max_tokens: int, api_key: str = "", base_url: str = "", system_instruction: str = "",
|
|
aspect_ratio: str = "auto", image_size: str = "auto", seed: int = 0) -> tuple:
|
|
"""生成图像"""
|
|
try:
|
|
if api_key.strip() or base_url.strip():
|
|
self.client.apply_runtime_config(api_key=api_key, base_url=base_url)
|
|
|
|
if not self.client.api_key:
|
|
return (None, "Error: API key not provided", "")
|
|
|
|
# 构建内容 - 请求图像生成
|
|
contents = [{
|
|
"role": "user",
|
|
"parts": [{"text": prompt}]
|
|
}]
|
|
|
|
# 生成配置 - 指定响应模式包含图像
|
|
generation_config = {
|
|
"temperature": temperature,
|
|
"topP": top_p,
|
|
"maxOutputTokens": max_tokens,
|
|
"responseModalities": ["TEXT", "IMAGE"], # 请求图像响应
|
|
}
|
|
|
|
# 添加图像配置
|
|
image_config = {}
|
|
if aspect_ratio != "auto":
|
|
image_config["aspectRatio"] = aspect_ratio
|
|
if image_size != "auto" and "gemini-2.5-flash-image" not in model:
|
|
image_config["imageSize"] = image_size
|
|
|
|
if image_config:
|
|
generation_config["imageConfig"] = image_config
|
|
|
|
# 添加seed
|
|
if seed > 0:
|
|
generation_config["seed"] = seed
|
|
|
|
# 构建响应信息
|
|
response_info = {
|
|
"model": model,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"max_tokens": max_tokens,
|
|
"aspect_ratio": aspect_ratio,
|
|
"image_size": image_size,
|
|
"seed": seed,
|
|
}
|
|
|
|
# 调用API
|
|
response = self.client.generate_content(
|
|
model=model,
|
|
contents=contents,
|
|
generation_config=generation_config,
|
|
system_instruction=system_instruction if system_instruction else None
|
|
)
|
|
|
|
# 提取图像和文本
|
|
if "candidates" in response and response["candidates"]:
|
|
parts = response["candidates"][0].get("content", {}).get("parts") or []
|
|
|
|
image_tensor = None
|
|
description = ""
|
|
|
|
for part in parts:
|
|
# 提取文本描述
|
|
if "text" in part:
|
|
description = part["text"]
|
|
|
|
# 提取图像数据
|
|
if "inlineData" in part:
|
|
inline_data = part["inlineData"]
|
|
if inline_data.get("mimeType", "").startswith("image/"):
|
|
base64_data = inline_data.get("data", "")
|
|
image_tensor = self._base64_to_tensor(base64_data)
|
|
|
|
# 添加生成信息到响应
|
|
response_info["status"] = "success"
|
|
response_info["image_generated"] = image_tensor is not None
|
|
response_info_str = json.dumps(response_info, ensure_ascii=False, indent=2)
|
|
|
|
if image_tensor is not None:
|
|
return (image_tensor, description, response_info_str)
|
|
else:
|
|
response_info["status"] = "error"
|
|
response_info["error"] = "No image generated"
|
|
response_info_str = json.dumps(response_info, ensure_ascii=False, indent=2)
|
|
return (None, f"Error: No image generated. Response: {description}", response_info_str)
|
|
else:
|
|
response_info["status"] = "error"
|
|
response_info["error"] = "No response from API"
|
|
response_info_str = json.dumps(response_info, ensure_ascii=False, indent=2)
|
|
return (None, "Error: No response from API", response_info_str)
|
|
|
|
except Exception as e:
|
|
response_info = {
|
|
"status": "error",
|
|
"error": str(e),
|
|
"model": model,
|
|
}
|
|
response_info_str = json.dumps(response_info, ensure_ascii=False, indent=2)
|
|
return (None, f"Error: {str(e)}", response_info_str)
|
|
|
|
|
|
class GoogleGeminiImageEditing:
|
|
"""Google Gemini 图像编辑节点"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"edit_prompt": ("STRING", {"multiline": True}),
|
|
"model": (["gemini-2.5-flash-image", "gemini-3-pro-image-preview", "gemini-3.1-flash-image-preview"], {"default": "gemini-2.5-flash-image"}),
|
|
},
|
|
"optional": {
|
|
"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
"image3": ("IMAGE",),
|
|
"image4": ("IMAGE",),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
"base_url": ("STRING", {"default": ""}),
|
|
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
|
|
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"max_tokens": ("INT", {"default": 8192, "min": 1, "max": 65536}),
|
|
"system_instruction": ("STRING", {"default": "", "multiline": True}),
|
|
"aspect_ratio": (["auto", "16:9", "4:3", "4:5", "3:2", "1:1", "2:3", "3:4", "5:4", "9:16", "21:9"], {"default": "auto"}),
|
|
"image_size": (["auto", "1K", "2K", "4K"], {"default": "auto"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 2147483647}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "STRING")
|
|
RETURN_NAMES = ("edited_image", "result", "response_info")
|
|
FUNCTION = "edit"
|
|
CATEGORY = "Google/Gemini"
|
|
|
|
def __init__(self):
|
|
self.client = GoogleGeminiClient()
|
|
|
|
def _tensor_to_base64(self, image_tensor: torch.Tensor) -> str:
|
|
"""将张量转换为base64"""
|
|
pil_images = tensor2pil(image_tensor)
|
|
pil_image = pil_images[0]
|
|
|
|
buffered = BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
return base64.b64encode(buffered.getvalue()).decode('utf-8')
|
|
|
|
def _base64_to_tensor(self, base64_str: str) -> torch.Tensor:
|
|
"""将base64转换为张量"""
|
|
try:
|
|
image_data = base64.b64decode(base64_str)
|
|
image = Image.open(BytesIO(image_data)).convert('RGB')
|
|
|
|
#save
|
|
buffered = BytesIO()
|
|
image.save(buffered, format="PNG")
|
|
|
|
return pil2tensor(image)
|
|
except Exception as e:
|
|
print(f"[GoogleAPI] Base64转换失败: {e}")
|
|
return None
|
|
|
|
def edit(self, edit_prompt: str, model: str, temperature: float, top_p: float,
|
|
max_tokens: int, api_key: str = "", base_url: str = "", system_instruction: str = "",
|
|
aspect_ratio: str = "auto", image_size: str = "auto", seed: int = 0,
|
|
image1=None, image2=None, image3=None, image4=None) -> tuple:
|
|
"""编辑多张图像"""
|
|
try:
|
|
if api_key.strip() or base_url.strip():
|
|
self.client.apply_runtime_config(api_key=api_key, base_url=base_url)
|
|
|
|
if not self.client.api_key:
|
|
return (None, "Error: API key not provided", "")
|
|
|
|
# 收集输入的图像
|
|
input_images = [image1, image2, image3, image4]
|
|
image_count = sum(1 for img in input_images if img is not None)
|
|
|
|
if image_count == 0:
|
|
return (None, "Error: No input images provided", "")
|
|
|
|
# 构建parts - 包含所有输入图像和编辑指令
|
|
parts = []
|
|
|
|
# 添加所有图像
|
|
for i, img in enumerate(input_images):
|
|
if img is not None:
|
|
image_base64 = self._tensor_to_base64(img)
|
|
parts.append({
|
|
"inlineData": {
|
|
"mimeType": "image/png",
|
|
"data": image_base64
|
|
}
|
|
})
|
|
|
|
# 添加编辑指令
|
|
parts.append({"text": edit_prompt})
|
|
|
|
# 构建内容
|
|
contents = [{
|
|
"role": "user",
|
|
"parts": parts
|
|
}]
|
|
|
|
# 生成配置 - 指定响应模式包含图像
|
|
generation_config = {
|
|
"temperature": temperature,
|
|
"topP": top_p,
|
|
"maxOutputTokens": max_tokens,
|
|
"responseModalities": ["TEXT", "IMAGE"], # 请求图像响应
|
|
}
|
|
|
|
# 添加图像配置
|
|
image_config = {}
|
|
if aspect_ratio != "auto":
|
|
image_config["aspectRatio"] = aspect_ratio
|
|
if image_size != "auto" and "gemini-2.5-flash-image" not in model:
|
|
image_config["imageSize"] = image_size
|
|
|
|
if image_config:
|
|
generation_config["imageConfig"] = image_config
|
|
|
|
# 添加seed
|
|
if seed > 0:
|
|
generation_config["seed"] = seed
|
|
|
|
# 构建响应信息
|
|
response_info = {
|
|
"model": model,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"max_tokens": max_tokens,
|
|
"aspect_ratio": aspect_ratio,
|
|
"image_size": image_size,
|
|
"seed": seed,
|
|
"input_image_count": image_count,
|
|
"edit_prompt": edit_prompt[:100] + "..." if len(edit_prompt) > 100 else edit_prompt,
|
|
}
|
|
|
|
# 调用API
|
|
response = self.client.generate_content(
|
|
model=model,
|
|
contents=contents,
|
|
generation_config=generation_config,
|
|
system_instruction=system_instruction if system_instruction else None
|
|
)
|
|
|
|
# 提取编辑后的图像和结果
|
|
if "candidates" in response and response["candidates"]:
|
|
parts = response["candidates"][0].get("content", {}).get("parts") or []
|
|
|
|
edited_image = None
|
|
result_text = ""
|
|
|
|
for part in parts:
|
|
# 提取文本结果
|
|
if "text" in part:
|
|
result_text = part["text"]
|
|
|
|
# 提取编辑后的图像
|
|
if "inlineData" in part:
|
|
inline_data = part["inlineData"]
|
|
print(f"[GoogleAPI] 检测到inlineData部分，mimeType={inline_data.get('mimeType', '')}")
|
|
if inline_data.get("mimeType", "").startswith("image/"):
|
|
base64_data = inline_data.get("data", "")
|
|
edited_image = self._base64_to_tensor(base64_data)
|
|
print(f"[GoogleAPI] 成功提取编辑后的图像")
|
|
|
|
#save_tensor_as_image(edited_image, "edited_image.png") # 调试用：保存编辑后的图像
|
|
|
|
|
|
# 添加编辑信息到响应
|
|
response_info["status"] = "success"
|
|
response_info["image_edited"] = edited_image is not None
|
|
response_info_str = json.dumps(response_info, ensure_ascii=False, indent=2)
|
|
|
|
if edited_image is not None:
|
|
return (edited_image, result_text, response_info_str)
|
|
else:
|
|
response_info["status"] = "error"
|
|
response_info["error"] = "No edited image generated"
|
|
response_info_str = json.dumps(response_info, ensure_ascii=False, indent=2)
|
|
return (None, f"Error: No edited image generated. Response: {result_text}", response_info_str)
|
|
else:
|
|
response_info["status"] = "error"
|
|
response_info["error"] = "No response from API"
|
|
response_info_str = json.dumps(response_info, ensure_ascii=False, indent=2)
|
|
return (None, "Error: No response from API", response_info_str)
|
|
|
|
except Exception as e:
|
|
response_info = {
|
|
"status": "error",
|
|
"error": str(e),
|
|
"model": model,
|
|
}
|
|
response_info_str = json.dumps(response_info, ensure_ascii=False, indent=2)
|
|
return (None, f"Error: {str(e)}", response_info_str)
|
|
|
|
|
|
class GoogleGeminiConfigManager:
|
|
"""Google Gemini 配置管理节点"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"api_key": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"action": (["set", "get", "clear"], {"default": "set"}),
|
|
"base_url": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("status",)
|
|
FUNCTION = "manage"
|
|
CATEGORY = "Google/Gemini"
|
|
|
|
def manage(self, api_key: str, action: str = "set", base_url: str = "") -> tuple:
|
|
"""管理配置"""
|
|
try:
|
|
if action == "set":
|
|
api_key_value = api_key.strip()
|
|
if not api_key_value:
|
|
return ("Error: API key cannot be empty",)
|
|
|
|
client = GoogleGeminiClient()
|
|
client.apply_runtime_config(api_key=api_key_value, base_url=base_url)
|
|
return (f"API key saved successfully (base_url: {describe_base_url(client.base_url)})",)
|
|
|
|
elif action == "get":
|
|
config = get_config()
|
|
api_key_value = config.get('api_key', '').strip()
|
|
if not api_key_value:
|
|
current_base_url = normalize_base_url(config.get("base_url", "")) or DEFAULT_BASE_URL
|
|
return (f"No API key configured; base_url: {describe_base_url(current_base_url)}",)
|
|
if api_key_value:
|
|
masked_key = api_key_value[:10] + "..." + api_key_value[-4:]
|
|
key_base_url_map = config.get("key_base_url_map", {})
|
|
mapped_base_url = ""
|
|
if isinstance(key_base_url_map, dict):
|
|
mapped_base_url = normalize_base_url(key_base_url_map.get(api_key_value, ""))
|
|
current_base_url = mapped_base_url or normalize_base_url(config.get("base_url", "")) or DEFAULT_BASE_URL
|
|
return (f"API key found: {masked_key}; base_url: {describe_base_url(current_base_url)}",)
|
|
|
|
elif action == "clear":
|
|
config = get_config()
|
|
current_key = config.get('api_key', '').strip()
|
|
key_base_url_map = config.get("key_base_url_map", {})
|
|
if isinstance(key_base_url_map, dict) and current_key in key_base_url_map:
|
|
key_base_url_map.pop(current_key, None)
|
|
config["key_base_url_map"] = key_base_url_map
|
|
config['api_key'] = ""
|
|
config['base_url'] = DEFAULT_BASE_URL
|
|
save_config(config)
|
|
return ("API key cleared; base_url reset to default",)
|
|
|
|
else:
|
|
return ("Error: Unknown action",)
|
|
|
|
except Exception as e:
|
|
return (f"Error: {str(e)}",)
|
|
|
|
|
|
# ==================== 节点注册 ====================
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GoogleGeminiImageAnalysis": GoogleGeminiImageAnalysis,
|
|
"GoogleGeminiImageGeneration": GoogleGeminiImageGeneration,
|
|
"GoogleGeminiImageEditing": GoogleGeminiImageEditing,
|
|
"GoogleGeminiConfigManager": GoogleGeminiConfigManager,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GoogleGeminiImageAnalysis": "Google Gemini - Image Analysis",
|
|
"GoogleGeminiImageGeneration": "Google Gemini - Image Generation",
|
|
"GoogleGeminiImageEditing": "Google Gemini - Image Editing",
|
|
"GoogleGeminiConfigManager": "Google Gemini - Config Manager",
|
|
}
|