1133 lines
46 KiB
Python
1133 lines
46 KiB
Python
"""
|
||
Gemini API 客户端
|
||
处理与 api.o1key.com 的通信，用于图像生成
|
||
"""
|
||
|
||
import re
|
||
import time
|
||
from io import BytesIO
|
||
from typing import Any, Callable, Dict, List, Optional
|
||
|
||
import aiohttp
|
||
from PIL import Image
|
||
|
||
from ..utils.image_utils import encode_image_to_base64, decode_base64_to_pil
|
||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||
from .base_client import BaseAPIClient
|
||
|
||
|
||
class GeminiAPIClient(BaseAPIClient):
|
||
"""
|
||
Gemini API 客户端
|
||
用于调用 Gemini 3 Pro 模型进行图像生成
|
||
"""
|
||
|
||
def __init__(self, api_key: Optional[str] = None):
|
||
"""
|
||
初始化客户端
|
||
|
||
Args:
|
||
api_key: API 密钥，如果为 None 则从配置文件或环境变量读取
|
||
"""
|
||
if api_key is None:
|
||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||
|
||
super().__init__(
|
||
base_url=get_api_base_url(),
|
||
api_key=api_key
|
||
)
|
||
|
||
@staticmethod
|
||
def build_proxy_url(port: str) -> Optional[str]:
|
||
"""
|
||
将端口号字符串转为 aiohttp 可用的 HTTP 代理 URL。
|
||
兼容 Windows / Mac，支持 v2rayN (10808) 和 Clash Verge (7897)。
|
||
|
||
Args:
|
||
port: 用户填写的端口号，如 "7897" 或 "10808"；空字符串返回 None
|
||
|
||
Returns:
|
||
代理 URL，如 "http://127.0.0.1:7897"，或 None（不使用代理）
|
||
"""
|
||
port = (port or "").strip()
|
||
if not port or not port.isdigit():
|
||
return None
|
||
return f"http://127.0.0.1:{port}"
|
||
|
||
def get_endpoint(self, model: str = "", resolution: str = "2K", image_format: str = "base64", **kwargs) -> str:
|
||
"""
|
||
根据模型和分辨率获取 API 端点
|
||
|
||
Args:
|
||
model: 模型名称
|
||
resolution: 分辨率（1K, 2K, 4K）
|
||
image_format: 返回格式，"url" 时追加 ?image_format=url 查询参数
|
||
|
||
Returns:
|
||
API 端点路径
|
||
"""
|
||
from ..models_config import get_model_endpoint
|
||
|
||
# 特殊处理：动态端点模型（根据分辨率选择）
|
||
if model == "nano-banana-pro-次卡":
|
||
if resolution == "1K":
|
||
endpoint = "/v1beta/models/nano-banana-pro:generateContent"
|
||
elif resolution == "2K":
|
||
endpoint = "/v1beta/models/nano-banana-pro-2k:generateContent"
|
||
elif resolution == "4K":
|
||
endpoint = "/v1beta/models/nano-banana-pro-4k:generateContent"
|
||
else:
|
||
endpoint = "/v1beta/models/nano-banana-pro-2k:generateContent"
|
||
|
||
elif model == "nano-banana-2-次卡":
|
||
if resolution == "512px":
|
||
endpoint = "/v1beta/models/nano-banana-2-0.5k:generateContent"
|
||
elif resolution == "1K":
|
||
endpoint = "/v1beta/models/nano-banana-2-1k:generateContent"
|
||
elif resolution == "2K":
|
||
endpoint = "/v1beta/models/nano-banana-2-2k:generateContent"
|
||
elif resolution == "4K":
|
||
endpoint = "/v1beta/models/nano-banana-2-4k:generateContent"
|
||
else:
|
||
endpoint = "/v1beta/models/nano-banana-2-2k:generateContent"
|
||
|
||
elif model == "nano-banana-2-官方计费":
|
||
if resolution == "512px":
|
||
endpoint = "/v1beta/models/nano-banana-2-0.5k-official:generateContent"
|
||
elif resolution == "1K":
|
||
endpoint = "/v1beta/models/nano-banana-2-1k-official:generateContent"
|
||
elif resolution == "2K":
|
||
endpoint = "/v1beta/models/nano-banana-2-2k-official:generateContent"
|
||
elif resolution == "4K":
|
||
endpoint = "/v1beta/models/nano-banana-2-4k-official:generateContent"
|
||
else:
|
||
endpoint = "/v1beta/models/nano-banana-2-2k-official:generateContent"
|
||
|
||
elif model == "nano-banana-pro-官方计费":
|
||
if resolution == "1K":
|
||
endpoint = "/v1beta/models/nano-banana-pro-1k-official:generateContent"
|
||
elif resolution == "2K":
|
||
endpoint = "/v1beta/models/nano-banana-pro-2k-official:generateContent"
|
||
elif resolution == "4K":
|
||
endpoint = "/v1beta/models/nano-banana-pro-4k-official:generateContent"
|
||
else:
|
||
endpoint = "/v1beta/models/nano-banana-pro-2k-official:generateContent"
|
||
|
||
elif model == "gemini-3-pro-image-preview-url":
|
||
if resolution == "1K":
|
||
endpoint = "/v1beta/models/gemini-3-pro-image-preview-url:generateContent"
|
||
elif resolution == "2K":
|
||
endpoint = "/v1beta/models/gemini-3-pro-image-preview-2k-url:generateContent"
|
||
elif resolution == "4K":
|
||
endpoint = "/v1beta/models/gemini-3-pro-image-preview-4k-url:generateContent"
|
||
else:
|
||
endpoint = "/v1beta/models/gemini-3-pro-image-preview-2k-url:generateContent"
|
||
|
||
else:
|
||
# 其他模型：从配置文件读取端点
|
||
endpoint = get_model_endpoint(model)
|
||
if not endpoint:
|
||
# 兜底：使用标准模式端点
|
||
endpoint = "/v1beta/models/gemini-3-pro-image-preview:generateContent"
|
||
|
||
# url 模式：追加查询参数
|
||
if image_format == "url":
|
||
endpoint = endpoint + "?image_format=url"
|
||
|
||
return endpoint
|
||
|
||
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
|
||
"""Gemini 请求 429/503/504 时返回中文错误文案。"""
|
||
if status_code == 429:
|
||
return "此型号资源暂时耗尽，继续重试即可"
|
||
if status_code == 503:
|
||
return "此型号目前需求量较大。需求高峰通常是暂时的。请稍后再试。"
|
||
if status_code == 504:
|
||
return "服务无法在截止期限内完成处理。可能原因是：您的提示词过大，无法及时处理。"
|
||
return None
|
||
|
||
@staticmethod
|
||
def _estimate_body_size(parts: list, extra_body: dict) -> int:
|
||
"""
|
||
快速估算请求体 JSON 序列化后的字节数。
|
||
extra_body 为除 contents[0].parts 以外的其余字段。
|
||
"""
|
||
import json
|
||
body = {
|
||
"contents": [{"role": "user", "parts": parts}],
|
||
**extra_body
|
||
}
|
||
return len(json.dumps(body).encode("utf-8"))
|
||
|
||
@staticmethod
|
||
def _scale_images_to_fit(
|
||
images: List[Image.Image],
|
||
scale: float
|
||
) -> List[Image.Image]:
|
||
"""
|
||
将所有图片按统一比例等比缩放（Lanczos，不降质量）。
|
||
scale < 1.0 时缩小，>= 1.0 时原样返回。
|
||
"""
|
||
if scale >= 1.0:
|
||
return images
|
||
result = []
|
||
for img in images:
|
||
new_w = max(1, int(img.width * scale))
|
||
new_h = max(1, int(img.height * scale))
|
||
result.append(img.resize((new_w, new_h), Image.Resampling.LANCZOS))
|
||
return result
|
||
|
||
def build_request_body(
|
||
self,
|
||
prompt: str = "",
|
||
images: Optional[List[Image.Image]] = None,
|
||
aspect_ratio: str = "1:1",
|
||
resolution: str = "2K",
|
||
enable_grounding: bool = False,
|
||
enable_image_search: bool = False,
|
||
image_compression: str = None,
|
||
thinking_level: str = None,
|
||
request_log_enabled: bool = True,
|
||
**kwargs
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
构建 API 请求体
|
||
|
||
Args:
|
||
prompt: 提示词
|
||
images: 输入图像列表（可选）
|
||
aspect_ratio: 宽高比
|
||
resolution: 分辨率
|
||
enable_grounding: 是否启用 Google Search Grounding
|
||
enable_image_search: 是否同时启用 Google Image Search（仅 Gemini 3.1 Flash 支持）
|
||
|
||
Returns:
|
||
请求体字典
|
||
"""
|
||
import json
|
||
|
||
_MAX_BODY_BYTES = 20 * 1024 * 1024 # 20 MB
|
||
_BODY_TARGET_BYTES = int(_MAX_BODY_BYTES * 0.9)
|
||
|
||
parts = []
|
||
|
||
# 添加文本部分
|
||
parts.append({"text": prompt})
|
||
|
||
image_config = {"imageSize": resolution}
|
||
if aspect_ratio and aspect_ratio != "智能":
|
||
image_config["aspectRatio"] = aspect_ratio
|
||
|
||
def _build_request_body(body_parts: List[dict]) -> Dict[str, Any]:
|
||
body = {
|
||
"contents": [
|
||
{
|
||
"role": "user",
|
||
"parts": body_parts
|
||
}
|
||
],
|
||
"generationConfig": {
|
||
"responseModalities": ["IMAGE"],
|
||
"imageConfig": image_config
|
||
}
|
||
}
|
||
|
||
# 添加思考深度配置
|
||
if thinking_level:
|
||
body["generationConfig"]["thinkingConfig"] = {
|
||
"thinkingLevel": thinking_level,
|
||
"includeThoughts": True
|
||
}
|
||
|
||
# 添加图片压缩参数
|
||
if image_compression:
|
||
body["image_compression"] = image_compression
|
||
|
||
# 添加 Google Search Grounding（如果启用）
|
||
# 新异步接口要求直接放在请求体顶层：{"google_search": true}
|
||
if enable_grounding or enable_image_search:
|
||
body["google_search"] = True
|
||
|
||
return body
|
||
|
||
def _request_size(body: Dict[str, Any]) -> int:
|
||
return len(json.dumps(body).encode("utf-8"))
|
||
|
||
def _format_size(size: int) -> str:
|
||
if size < 1024 * 1024:
|
||
return f"{size / 1024:.2f}KB"
|
||
return f"{size / 1024 / 1024:.2f}MB"
|
||
|
||
def _shorten_base64_for_log(obj, max_len: int = 200):
|
||
if isinstance(obj, dict):
|
||
result = {}
|
||
for key, value in obj.items():
|
||
if key == "data" and isinstance(value, str) and len(value) > max_len:
|
||
result[key] = f"<base64 data, {len(value)} chars>"
|
||
else:
|
||
result[key] = _shorten_base64_for_log(value, max_len)
|
||
return result
|
||
if isinstance(obj, list):
|
||
return [_shorten_base64_for_log(item, max_len) for item in obj]
|
||
return obj
|
||
|
||
def _log_original_request_body(body: Dict[str, Any]) -> None:
|
||
body_size = _request_size(body)
|
||
print(
|
||
f"\n{'=' * 60}\n"
|
||
f"[原始请求体日志] 请求体积: {_format_size(body_size)} "
|
||
f"(inline_data.data 已折叠显示 base64 长度)\n"
|
||
f"{json.dumps(_shorten_base64_for_log(body), ensure_ascii=False, indent=2)}\n"
|
||
f"{'=' * 60}\n"
|
||
)
|
||
|
||
original_request_logged = False
|
||
|
||
# 添加图像部分（如果有）
|
||
if images:
|
||
def _build_image_parts(src_images: List[Image.Image]) -> List[dict]:
|
||
img_parts = []
|
||
for img in src_images:
|
||
img_base64 = encode_image_to_base64(img)
|
||
img_parts.append({
|
||
"inline_data": {
|
||
"mime_type": "image/png",
|
||
"data": img_base64
|
||
}
|
||
})
|
||
return img_parts
|
||
|
||
def _estimate_with_images(img_parts: List[dict]) -> int:
|
||
return _request_size(_build_request_body(parts + img_parts))
|
||
|
||
working_images = list(images)
|
||
img_parts = _build_image_parts(working_images)
|
||
original_request_body = _build_request_body(parts + img_parts)
|
||
if request_log_enabled:
|
||
_log_original_request_body(original_request_body)
|
||
original_request_logged = True
|
||
estimated = _estimate_with_images(img_parts)
|
||
|
||
if estimated > _MAX_BODY_BYTES:
|
||
ratio = _BODY_TARGET_BYTES / estimated
|
||
scale = ratio ** 0.5 # 面积比 → 线性比
|
||
working_images = self._scale_images_to_fit(working_images, scale)
|
||
img_parts = _build_image_parts(working_images)
|
||
estimated = _estimate_with_images(img_parts)
|
||
|
||
orig_sizes = ", ".join(f"{img.width}×{img.height}" for img in images)
|
||
new_sizes = ", ".join(f"{img.width}×{img.height}" for img in working_images)
|
||
size_mb = estimated / (1024 * 1024)
|
||
target_mb = _BODY_TARGET_BYTES / (1024 * 1024)
|
||
print(
|
||
f"Nano Banana Pro: 输入图片已按请求体目标大小自动缩放\n"
|
||
f" 原始尺寸: {orig_sizes}\n"
|
||
f" 缩放后: {new_sizes}\n"
|
||
f" 请求体积: {size_mb:.2f}MB（目标 {target_mb:.2f}MB，限制 20MB）"
|
||
)
|
||
|
||
parts.extend(img_parts)
|
||
|
||
request_body = _build_request_body(parts)
|
||
if request_log_enabled and not original_request_logged:
|
||
_log_original_request_body(request_body)
|
||
|
||
request_size = _request_size(request_body)
|
||
if request_size > _MAX_BODY_BYTES:
|
||
raise ValueError(
|
||
f"请求体超过 20MB 限制（当前 {request_size / 1024 / 1024:.2f}MB），"
|
||
"已停止提交；请减少参考图数量、降低图片复杂度或缩短提示词"
|
||
)
|
||
|
||
return request_body
|
||
|
||
def parse_response(self, response: Dict[str, Any]) -> List[Image.Image]:
|
||
"""
|
||
同步解析 API 响应（保留以满足抽象基类要求）
|
||
|
||
注意：此方法仅用于兼容基类接口，实际使用请调用 parse_response_async()
|
||
|
||
Args:
|
||
response: API 响应字典
|
||
|
||
Returns:
|
||
图像列表
|
||
|
||
Raises:
|
||
RuntimeError: 此方法不应被直接调用
|
||
"""
|
||
raise RuntimeError(
|
||
"parse_response() 不应被直接调用。"
|
||
"请使用 generate_single_async() 或 generate_batch_async() 等高级方法。"
|
||
)
|
||
|
||
async def parse_response_async(
|
||
self,
|
||
response: Dict[str, Any],
|
||
session: Optional[aiohttp.ClientSession] = None
|
||
) -> tuple[List[Image.Image], Dict[str, Any]]:
|
||
"""
|
||
异步解析 API 响应，提取生成的图像
|
||
|
||
Args:
|
||
response: API 响应字典
|
||
session: aiohttp 会话（用于下载图片）
|
||
|
||
Returns:
|
||
(图像列表, 格式信息字典)
|
||
格式信息包含: type (base64/url), size, resolution, download_speed (仅URL)
|
||
|
||
Raises:
|
||
RuntimeError: 解析失败或 API 拒绝时
|
||
"""
|
||
|
||
# 初始化格式信息
|
||
format_info = {
|
||
"type": None, # "base64" or "url"
|
||
"size": 0,
|
||
"resolution": None,
|
||
"download_speed": None
|
||
}
|
||
|
||
candidates = response.get("candidates", [])
|
||
|
||
# ========== 错误检测（按优先级顺序）==========
|
||
|
||
# 1. 检查 candidatesTokenCount（最高优先级）
|
||
usage_metadata = response.get("usageMetadata", {})
|
||
candidates_token_count = usage_metadata.get("candidatesTokenCount", -1)
|
||
|
||
if candidates_token_count == 0:
|
||
error_msg = (
|
||
"Damn!你触发顶级风控啦!还没到生图阶段就被拒了。\n"
|
||
"赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~"
|
||
)
|
||
raise RuntimeError(error_msg)
|
||
|
||
# 2. 检查 finishReason（次优先级）
|
||
candidates = response.get("candidates", [])
|
||
if candidates:
|
||
for candidate in candidates:
|
||
finish_reason = candidate.get("finishReason", "")
|
||
|
||
if finish_reason and finish_reason != "STOP":
|
||
error_msg = (
|
||
"Ohh no! 生图过程触发风控,图片被拒绝生成!\n"
|
||
"可能原因如下:\n"
|
||
"1.违禁内容\n"
|
||
"2.触发安全过滤器\n"
|
||
"3.涉及版权问题\n"
|
||
"4. Token超限\n"
|
||
"赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~"
|
||
)
|
||
raise RuntimeError(error_msg)
|
||
|
||
# ========== 图像提取 ==========
|
||
|
||
images = []
|
||
text_responses = [] # 收集文本响应
|
||
|
||
# 需要关闭 session 的标记
|
||
close_session = False
|
||
if session is None:
|
||
session = self._make_session()
|
||
close_session = True
|
||
|
||
try:
|
||
for candidate_idx, candidate in enumerate(candidates):
|
||
content = candidate.get("content", {})
|
||
parts = content.get("parts", [])
|
||
|
||
for part_idx, part in enumerate(parts):
|
||
# 方式1: inline_data 或 inlineData (base64)
|
||
# 兼容两种命名方式：蛇形（inline_data）和驼峰（inlineData）
|
||
inline_data_key = None
|
||
if "inline_data" in part:
|
||
inline_data_key = "inline_data"
|
||
elif "inlineData" in part:
|
||
inline_data_key = "inlineData"
|
||
|
||
if inline_data_key:
|
||
# 跳过思考链草稿图（thought:true 标记的 inlineData 为模型自检用途，非最终输出）
|
||
if part.get("thought") is True:
|
||
continue
|
||
|
||
inline_data = part[inline_data_key]
|
||
# 同样兼容 data/mimeType 的命名
|
||
img_data = inline_data.get("data") or inline_data.get("data", "")
|
||
|
||
if img_data:
|
||
img = decode_base64_to_pil(img_data)
|
||
images.append(img)
|
||
|
||
# 记录格式信息
|
||
if format_info["type"] is None:
|
||
format_info["type"] = "base64"
|
||
format_info["size"] = len(img_data) * 3 / 4 # Base64 解码后的字节数
|
||
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
|
||
|
||
# 方式1b: fileData（?image_format=url 模式下服务端返回 URL 替代 inlineData）
|
||
elif "fileData" in part:
|
||
if part.get("thought") is True:
|
||
continue
|
||
file_data = part["fileData"]
|
||
url = file_data.get("fileUri") or file_data.get("file_uri", "")
|
||
if url:
|
||
try:
|
||
download_start = time.time()
|
||
async with session.get(url) as img_response:
|
||
if img_response.status == 200:
|
||
img_bytes = await img_response.read()
|
||
download_time = time.time() - download_start
|
||
img_size = len(img_bytes)
|
||
speed = img_size / download_time if download_time > 0 else 0
|
||
|
||
img = Image.open(BytesIO(img_bytes))
|
||
images.append(img)
|
||
|
||
if format_info["type"] is None:
|
||
format_info["type"] = "url"
|
||
format_info["size"] = img_size
|
||
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
|
||
format_info["download_speed"] = speed
|
||
except Exception:
|
||
pass # 静默失败
|
||
|
||
# 方式2: text 中的 URL - 改为异步下载
|
||
elif "text" in part:
|
||
text = part["text"]
|
||
|
||
# 收集文本响应（用于后续错误检测）
|
||
text_responses.append(text)
|
||
|
||
# 尝试 markdown 格式: 
|
||
url_pattern_md = r'!\[.*?\]\((https?://[^\)]+)\)'
|
||
urls = re.findall(url_pattern_md, text)
|
||
|
||
# 如果没找到，尝试纯 URL 格式
|
||
if not urls:
|
||
url_pattern_plain = r'https?://[^\s<>"{}|\\^`\[\]]+'
|
||
urls = re.findall(url_pattern_plain, text)
|
||
|
||
if urls:
|
||
for url_idx, url in enumerate(urls):
|
||
try:
|
||
# 使用 aiohttp 异步下载
|
||
download_start = time.time()
|
||
async with session.get(url) as img_response:
|
||
if img_response.status == 200:
|
||
img_data = await img_response.read()
|
||
download_time = time.time() - download_start
|
||
img_size = len(img_data)
|
||
speed = img_size / download_time if download_time > 0 else 0
|
||
|
||
img = Image.open(BytesIO(img_data))
|
||
images.append(img)
|
||
|
||
# 记录格式信息（只记录第一张）
|
||
if format_info["type"] is None:
|
||
format_info["type"] = "url"
|
||
format_info["size"] = img_size
|
||
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
|
||
format_info["download_speed"] = speed
|
||
except Exception as e:
|
||
pass # 静默失败，继续尝试其他URL
|
||
|
||
# 方式3: 直接的 URL 字段 - 也改为异步
|
||
elif "imageUrl" in part or "url" in part:
|
||
url = part.get("imageUrl") or part.get("url")
|
||
try:
|
||
download_start = time.time()
|
||
async with session.get(url) as img_response:
|
||
if img_response.status == 200:
|
||
img_data = await img_response.read()
|
||
download_time = time.time() - download_start
|
||
img_size = len(img_data)
|
||
speed = img_size / download_time if download_time > 0 else 0
|
||
|
||
img = Image.open(BytesIO(img_data))
|
||
images.append(img)
|
||
|
||
# 记录格式信息
|
||
if format_info["type"] is None:
|
||
format_info["type"] = "url"
|
||
format_info["size"] = img_size
|
||
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
|
||
format_info["download_speed"] = speed
|
||
except Exception as e:
|
||
pass # 静默失败
|
||
|
||
except Exception as e:
|
||
raise RuntimeError(f"解析 API 响应失败: {str(e)}")
|
||
|
||
finally:
|
||
if close_session:
|
||
await session.close()
|
||
|
||
# 3. 检查 API 文本响应拒绝说明
|
||
if not images and text_responses:
|
||
# API 返回了文本但没有图片，说明请求被拒绝
|
||
combined_text = "\n".join(text_responses)
|
||
error_msg = (
|
||
f"API 拒绝响应\n\n"
|
||
f"API 返回说明：\n{combined_text}\n\n"
|
||
f"建议：\n"
|
||
f" - 根据上述说明调整请求内容\n"
|
||
f" - 确保提示词和参考图符合使用规范"
|
||
)
|
||
raise RuntimeError(error_msg)
|
||
|
||
if not images:
|
||
raise RuntimeError("API 响应中未找到生成的图像")
|
||
|
||
return images, format_info
|
||
|
||
async def generate_single_async(
|
||
self,
|
||
prompt: str,
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
images: Optional[List[Image.Image]] = None,
|
||
session=None,
|
||
task_index: Optional[int] = None,
|
||
total_tasks: Optional[int] = None,
|
||
debug: bool = False,
|
||
debug_request: bool = False,
|
||
enable_grounding: bool = False,
|
||
enable_image_search: bool = False,
|
||
image_format: str = "base64",
|
||
thinking_level: str = None,
|
||
) -> tuple[List[Image.Image], Dict[str, Any]]:
|
||
"""
|
||
单次异步生成请求（极简单行日志）
|
||
|
||
Args:
|
||
prompt: 提示词
|
||
model: 模型名称
|
||
resolution: 分辨率
|
||
aspect_ratio: 宽高比
|
||
images: 输入图像列表
|
||
session: aiohttp 会话
|
||
task_index: 任务索引（用于批量任务）
|
||
total_tasks: 总任务数（用于批量任务）
|
||
debug: 是否打印完整 API 响应
|
||
debug_request: 是否打印发送的请求体（base64 图片数据将被截断）
|
||
enable_grounding: 是否启用 Google Search Grounding
|
||
enable_image_search: 是否同时启用 Google Image Search
|
||
|
||
Returns:
|
||
(生成的图像列表, 计时信息字典)
|
||
"""
|
||
import json
|
||
|
||
total_start = time.time()
|
||
|
||
# 任务前缀
|
||
task_prefix = f"[{task_index}/{total_tasks}] " if task_index is not None and total_tasks else ""
|
||
|
||
# ========== 1. 构建请求 ==========
|
||
build_start = time.time()
|
||
endpoint = self.get_endpoint(model=model, resolution=resolution, image_format=image_format)
|
||
request_body = self.build_request_body(
|
||
prompt=prompt,
|
||
images=images,
|
||
aspect_ratio=aspect_ratio,
|
||
resolution=resolution,
|
||
enable_grounding=enable_grounding,
|
||
enable_image_search=enable_image_search,
|
||
thinking_level=thinking_level,
|
||
)
|
||
build_time = time.time() - build_start
|
||
|
||
# ========== 调试日志：打印请求体 ==========
|
||
if debug_request:
|
||
import json as _json
|
||
|
||
def _truncate_base64_req(obj, max_len=200):
|
||
if isinstance(obj, dict):
|
||
return {k: _truncate_base64_req(v, max_len) for k, v in obj.items()}
|
||
elif isinstance(obj, list):
|
||
return [_truncate_base64_req(item, max_len) for item in obj]
|
||
elif isinstance(obj, str) and len(obj) > max_len:
|
||
if all(c in 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=' for c in obj[:50]):
|
||
return f"<base64 data, {len(obj)} chars>"
|
||
return obj
|
||
return obj
|
||
|
||
safe_request = _truncate_base64_req(request_body)
|
||
print(
|
||
f"\n{'='*60}\n"
|
||
f"[请求体日志] {task_prefix}发送请求体：\n"
|
||
f"端点: {endpoint}\n"
|
||
f"{_json.dumps(safe_request, ensure_ascii=False, indent=2)}\n"
|
||
f"{'='*60}\n"
|
||
)
|
||
|
||
# 计算请求体大小
|
||
request_size = len(json.dumps(request_body).encode('utf-8'))
|
||
if request_size < 1024 * 1024:
|
||
size_str = f"{request_size / 1024:.2f}KB"
|
||
else:
|
||
size_str = f"{request_size / (1024 * 1024):.2f}MB"
|
||
|
||
# ========== 2. 发送网络请求 ==========
|
||
request_start = time.time()
|
||
|
||
try:
|
||
response = await self.request_async(endpoint, request_body, session)
|
||
except Exception as e:
|
||
raise
|
||
|
||
request_time = time.time() - request_start
|
||
|
||
# ========== 调试日志：打印完整响应 ==========
|
||
if debug:
|
||
import json as _json
|
||
# 构建可安全序列化的响应副本（截断 base64 图片数据避免输出过长）
|
||
def _truncate_base64(obj, max_len=200):
|
||
if isinstance(obj, dict):
|
||
return {k: _truncate_base64(v, max_len) for k, v in obj.items()}
|
||
elif isinstance(obj, list):
|
||
return [_truncate_base64(item, max_len) for item in obj]
|
||
elif isinstance(obj, str) and len(obj) > max_len:
|
||
# 判断是否为 base64 图片数据（不含空格/换行的长字符串）
|
||
if all(c in 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=' for c in obj[:50]):
|
||
return f"<base64 data, {len(obj)} chars>"
|
||
return obj
|
||
return obj
|
||
|
||
safe_response = _truncate_base64(response)
|
||
print(
|
||
f"\n{'='*60}\n"
|
||
f"[调试日志] {task_prefix}完整 API 响应：\n"
|
||
f"{_json.dumps(safe_response, ensure_ascii=False, indent=2)}\n"
|
||
f"{'='*60}\n"
|
||
)
|
||
|
||
# ========== 3. 解析响应 ==========
|
||
parse_start = time.time()
|
||
|
||
try:
|
||
result_images, format_info = await self.parse_response_async(response, session)
|
||
except Exception as e:
|
||
parse_time = time.time() - parse_start
|
||
error_first_line = str(e).split('\n')[0]
|
||
print(f"{task_prefix}请求 {size_str} → API {request_time:.1f}s → 解析失败: {error_first_line} ✗")
|
||
raise
|
||
|
||
parse_time = time.time() - parse_start
|
||
|
||
# ========== 4. 格式化输出（单行） ==========
|
||
# 格式化图像大小
|
||
img_size = format_info.get("size", 0)
|
||
if img_size < 1024 * 1024:
|
||
img_size_str = f"{img_size / 1024:.2f}KB"
|
||
else:
|
||
img_size_str = f"{img_size / (1024 * 1024):.2f}MB"
|
||
|
||
# 根据类型构建下载信息
|
||
if format_info.get("type") == "base64":
|
||
download_info = f"Base64 {img_size_str} ({parse_time:.1f}s)"
|
||
elif format_info.get("type") == "url":
|
||
speed = format_info.get("download_speed", 0)
|
||
speed_str = f"{speed / (1024 * 1024):.1f}MB/s"
|
||
download_info = f"URL {img_size_str} ({parse_time:.1f}s, {speed_str})"
|
||
else:
|
||
download_info = f"{img_size_str}"
|
||
|
||
# 单行输出
|
||
timing = response.get("_timing", {})
|
||
net_connect = timing.get("connect_time")
|
||
net_download = timing.get("download_time")
|
||
if net_connect is not None and net_download is not None:
|
||
net_str = f" | 连接 {net_connect:.2f}s | 下载 {net_download:.2f}s"
|
||
else:
|
||
net_str = ""
|
||
print(f"{task_prefix}请求 {size_str} → API {request_time:.1f}s → {download_info} ✓{net_str}")
|
||
|
||
# 返回结果和计时信息
|
||
total_time = time.time() - total_start
|
||
timing_info = {
|
||
"build_time": build_time,
|
||
"request_time": request_time,
|
||
"parse_time": parse_time,
|
||
"total_time": total_time,
|
||
"format_type": format_info.get("type", "unknown")
|
||
}
|
||
|
||
return result_images, timing_info
|
||
|
||
async def generate_batch_async(
|
||
self,
|
||
prompt: str,
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
batch_size: int,
|
||
images: Optional[List[Image.Image]] = None,
|
||
progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None,
|
||
debug: bool = False,
|
||
debug_request: bool = False,
|
||
enable_grounding: bool = False,
|
||
enable_image_search: bool = False,
|
||
image_format: str = "base64",
|
||
thinking_level: str = None,
|
||
) -> List[Image.Image]:
|
||
"""
|
||
批量全并发生成 - 改进版：支持分批处理和内存管理
|
||
|
||
Args:
|
||
prompt: 提示词
|
||
model: 模型名称
|
||
resolution: 分辨率
|
||
aspect_ratio: 宽高比
|
||
batch_size: 批次大小
|
||
images: 输入图像列表
|
||
progress_callback: 进度回调，签名为 (completed, total, success, error_msg)
|
||
debug: 是否打印完整 API 响应
|
||
debug_request: 是否打印发送的请求体
|
||
enable_grounding: 是否启用 Google Search Grounding
|
||
enable_image_search: 是否同时启用 Google Image Search
|
||
|
||
Returns:
|
||
生成的图像列表
|
||
"""
|
||
import aiohttp
|
||
import asyncio
|
||
|
||
all_images = []
|
||
completed = 0
|
||
success_count = 0
|
||
fail_count = 0
|
||
first_error = None # 保存第一个错误
|
||
|
||
# 分批处理配置
|
||
max_concurrent = 10 # 最大并发数
|
||
save_batch_size = 10 # 分批保存大小
|
||
|
||
# 计算需要多少批次
|
||
num_batches = (batch_size + max_concurrent - 1) // max_concurrent
|
||
|
||
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
|
||
|
||
async with aiohttp.ClientSession(connector=connector) as session:
|
||
# 分批执行
|
||
for batch_idx in range(num_batches):
|
||
batch_start = batch_idx * max_concurrent
|
||
batch_end = min(batch_start + max_concurrent, batch_size)
|
||
batch_size_current = batch_end - batch_start
|
||
|
||
# 创建当前批次的任务
|
||
tasks = []
|
||
for i in range(batch_size_current):
|
||
task_index = batch_start + i
|
||
task = asyncio.create_task(
|
||
self.generate_single_async(
|
||
prompt=prompt,
|
||
model=model,
|
||
resolution=resolution,
|
||
aspect_ratio=aspect_ratio,
|
||
images=images,
|
||
session=session,
|
||
task_index=task_index + 1,
|
||
total_tasks=batch_size,
|
||
debug=debug,
|
||
debug_request=debug_request,
|
||
enable_grounding=enable_grounding,
|
||
enable_image_search=enable_image_search,
|
||
image_format=image_format,
|
||
thinking_level=thinking_level,
|
||
),
|
||
name=f"task_{task_index}"
|
||
)
|
||
tasks.append(task)
|
||
|
||
# 收集当前批次的结果
|
||
batch_images = []
|
||
batch_completed = 0
|
||
|
||
for coro in asyncio.as_completed(tasks):
|
||
batch_completed += 1
|
||
completed += 1
|
||
|
||
try:
|
||
result_images, timing_info = await coro
|
||
if result_images:
|
||
# 立即处理生成的图片
|
||
for img in result_images:
|
||
batch_images.append(img)
|
||
all_images.append(img)
|
||
|
||
success_count += 1
|
||
|
||
# 通知进度
|
||
if progress_callback:
|
||
progress_callback(completed, batch_size, True, None)
|
||
|
||
except Exception as e:
|
||
fail_count += 1
|
||
# 保存第一个错误（用于后续抛出）
|
||
if first_error is None:
|
||
first_error = e
|
||
error_msg = str(e)
|
||
|
||
# 传递完整的错误信息（用于排查问题）
|
||
if progress_callback:
|
||
progress_callback(completed, batch_size, False, error_msg)
|
||
|
||
# 当前批次完成后，立即清理内存
|
||
if batch_images:
|
||
import gc
|
||
gc.collect()
|
||
await asyncio.sleep(0.1)
|
||
|
||
# 清空当前批次图片引用，帮助垃圾回收
|
||
batch_images = []
|
||
|
||
# 最终结果检查
|
||
if not all_images:
|
||
# 如果有保存的原始错误，直接抛出原始错误
|
||
if first_error:
|
||
raise first_error
|
||
raise RuntimeError(f"批量生成失败，{fail_count} 个请求全部失败")
|
||
|
||
if batch_size > 1:
|
||
print(f"GeminiClient: 批量生成完成，成功 {success_count}/{batch_size}，失败 {fail_count}")
|
||
return all_images
|
||
|
||
def generate_sync(
|
||
self,
|
||
prompt: str,
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
batch_size: int,
|
||
images: Optional[List[Image.Image]] = None,
|
||
progress_callback: Optional[Callable[[int, int], None]] = None,
|
||
debug: bool = False,
|
||
debug_request: bool = False,
|
||
enable_grounding: bool = False,
|
||
enable_image_search: bool = False,
|
||
image_format: str = "base64",
|
||
thinking_level: str = None,
|
||
) -> List[Image.Image]:
|
||
"""
|
||
同步生成接口（用于 ComfyUI）
|
||
|
||
Args:
|
||
prompt: 提示词
|
||
model: 模型名称
|
||
resolution: 分辨率
|
||
aspect_ratio: 宽高比
|
||
batch_size: 批次大小
|
||
images: 输入图像列表
|
||
progress_callback: 进度回调
|
||
debug: 是否打印完整 API 响应
|
||
debug_request: 是否打印发送的请求体
|
||
enable_grounding: 是否启用 Google Search Grounding
|
||
enable_image_search: 是否同时启用 Google Image Search
|
||
|
||
Returns:
|
||
生成的图像列表
|
||
"""
|
||
coro = self.generate_batch_async(
|
||
prompt=prompt,
|
||
model=model,
|
||
resolution=resolution,
|
||
aspect_ratio=aspect_ratio,
|
||
batch_size=batch_size,
|
||
images=images,
|
||
progress_callback=progress_callback,
|
||
debug=debug,
|
||
debug_request=debug_request,
|
||
enable_grounding=enable_grounding,
|
||
enable_image_search=enable_image_search,
|
||
image_format=image_format,
|
||
thinking_level=thinking_level,
|
||
)
|
||
|
||
return self.run_async_in_thread(coro)
|
||
|
||
async def generate_multi_prompts_async(
|
||
self,
|
||
prompts: List[str],
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
images_per_prompt: int,
|
||
images: Optional[List[Image.Image]] = None,
|
||
progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None,
|
||
debug: bool = False,
|
||
debug_request: bool = False,
|
||
enable_grounding: bool = False,
|
||
enable_image_search: bool = False
|
||
) -> List[Image.Image]:
|
||
"""
|
||
多提示词批量生成 - 改进版：支持分批处理和内存管理
|
||
|
||
为每个提示词生成指定数量的图像，分批并发执行。
|
||
|
||
Args:
|
||
prompts: 提示词列表
|
||
model: 模型名称
|
||
resolution: 分辨率
|
||
aspect_ratio: 宽高比
|
||
images_per_prompt: 每个提示词生成的图像数量
|
||
images: 输入图像列表（所有提示词共享）
|
||
progress_callback: 进度回调，签名为 (completed, total, success, error_msg)
|
||
debug: 是否打印完整 API 响应
|
||
debug_request: 是否打印发送的请求体
|
||
enable_grounding: 是否启用 Google Search Grounding
|
||
enable_image_search: 是否同时启用 Google Image Search
|
||
|
||
Returns:
|
||
生成的图像列表（长度 = len(prompts) * images_per_prompt）
|
||
"""
|
||
import aiohttp
|
||
import asyncio
|
||
|
||
all_images = []
|
||
completed = 0
|
||
success_count = 0
|
||
fail_count = 0
|
||
first_error = None # 保存第一个错误
|
||
total_tasks = len(prompts) * images_per_prompt
|
||
|
||
# 分批处理配置
|
||
max_concurrent = 10 # 最大并发数
|
||
|
||
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
|
||
|
||
async with aiohttp.ClientSession(connector=connector) as session:
|
||
# 创建所有任务
|
||
tasks = []
|
||
task_idx = 0
|
||
for prompt in prompts:
|
||
for _ in range(images_per_prompt):
|
||
task = asyncio.create_task(
|
||
self.generate_single_async(
|
||
prompt=prompt,
|
||
model=model,
|
||
resolution=resolution,
|
||
aspect_ratio=aspect_ratio,
|
||
images=images,
|
||
session=session,
|
||
task_index=task_idx + 1,
|
||
total_tasks=total_tasks,
|
||
debug=debug,
|
||
debug_request=debug_request,
|
||
enable_grounding=enable_grounding,
|
||
enable_image_search=enable_image_search
|
||
),
|
||
name=f"task_{task_idx}"
|
||
)
|
||
tasks.append(task)
|
||
task_idx += 1
|
||
|
||
# 分批处理：每10个任务为一组
|
||
batch_size = max_concurrent
|
||
num_batches = (total_tasks + batch_size - 1) // batch_size
|
||
|
||
for batch_idx in range(num_batches):
|
||
batch_start = batch_idx * batch_size
|
||
batch_end = min(batch_start + batch_size, total_tasks)
|
||
batch_tasks = tasks[batch_start:batch_end]
|
||
|
||
# 收集当前批次的结果
|
||
batch_images = []
|
||
|
||
for coro in asyncio.as_completed(batch_tasks):
|
||
completed += 1
|
||
|
||
try:
|
||
result_images, timing_info = await coro
|
||
if result_images:
|
||
# 立即处理生成的图片
|
||
for img in result_images:
|
||
batch_images.append(img)
|
||
all_images.append(img)
|
||
|
||
success_count += 1
|
||
|
||
# 通知进度
|
||
if progress_callback:
|
||
progress_callback(completed, total_tasks, True, None)
|
||
|
||
except Exception as e:
|
||
fail_count += 1
|
||
# 保存第一个错误（用于后续抛出）
|
||
if first_error is None:
|
||
first_error = e
|
||
error_msg = str(e)
|
||
|
||
# 传递完整的错误信息（用于排查问题）
|
||
if progress_callback:
|
||
progress_callback(completed, total_tasks, False, error_msg)
|
||
|
||
# 当前批次完成后，立即清理内存
|
||
if batch_images:
|
||
import gc
|
||
gc.collect()
|
||
await asyncio.sleep(0.1)
|
||
|
||
# 清空当前批次图片引用，帮助垃圾回收
|
||
batch_images = []
|
||
|
||
if not all_images:
|
||
# 如果有保存的原始错误，直接抛出原始错误
|
||
if first_error:
|
||
raise first_error
|
||
raise RuntimeError(f"批量生成失败，{fail_count} 个请求全部失败")
|
||
|
||
return all_images
|
||
|
||
def generate_multi_prompts_sync(
|
||
self,
|
||
prompts: List[str],
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
images_per_prompt: int,
|
||
images: Optional[List[Image.Image]] = None,
|
||
progress_callback: Optional[Callable[[int, int], None]] = None,
|
||
debug: bool = False,
|
||
debug_request: bool = False,
|
||
enable_grounding: bool = False,
|
||
enable_image_search: bool = False
|
||
) -> List[Image.Image]:
|
||
"""
|
||
多提示词批量生成（同步接口，用于 ComfyUI）
|
||
|
||
Args:
|
||
prompts: 提示词列表
|
||
model: 模型名称
|
||
resolution: 分辨率
|
||
aspect_ratio: 宽高比
|
||
images_per_prompt: 每个提示词生成的图像数量
|
||
images: 输入图像列表
|
||
progress_callback: 进度回调
|
||
debug: 是否打印完整 API 响应
|
||
debug_request: 是否打印发送的请求体
|
||
enable_grounding: 是否启用 Google Search Grounding
|
||
enable_image_search: 是否同时启用 Google Image Search
|
||
|
||
Returns:
|
||
生成的图像列表
|
||
"""
|
||
coro = self.generate_multi_prompts_async(
|
||
prompts=prompts,
|
||
model=model,
|
||
resolution=resolution,
|
||
aspect_ratio=aspect_ratio,
|
||
images_per_prompt=images_per_prompt,
|
||
images=images,
|
||
progress_callback=progress_callback,
|
||
debug=debug,
|
||
debug_request=debug_request,
|
||
enable_grounding=enable_grounding,
|
||
enable_image_search=enable_image_search
|
||
)
|
||
|
||
return self.run_async_in_thread(coro)
|