Publish current ComfyUI O1Key code baseline
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
This commit is contained in:
Jony
+68
-50
@@ -1,12 +1,13 @@
|
||||
"""
|
||||
Gemini API 客户端
|
||||
处理与 api.o1key.com 的通信，用于图像生成
|
||||
处理与 api.o1key.cn 的通信，用于图像生成
|
||||
"""
|
||||
|
||||
import base64
|
||||
import re
|
||||
import time
|
||||
from io import BytesIO
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
from PIL import Image
|
||||
@@ -80,7 +81,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
endpoint = "/v1beta/models/nano-banana-pro-2k:generateContent"
|
||||
|
||||
elif model == "nano-banana-2-次卡":
|
||||
if resolution == "512px":
|
||||
if resolution == "512":
|
||||
endpoint = "/v1beta/models/nano-banana-2-0.5k:generateContent"
|
||||
elif resolution == "1K":
|
||||
endpoint = "/v1beta/models/nano-banana-2-1k:generateContent"
|
||||
@@ -92,7 +93,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
endpoint = "/v1beta/models/nano-banana-2-2k:generateContent"
|
||||
|
||||
elif model == "nano-banana-2-官方计费":
|
||||
if resolution == "512px":
|
||||
if resolution == "512":
|
||||
endpoint = "/v1beta/models/nano-banana-2-0.5k-official:generateContent"
|
||||
elif resolution == "1K":
|
||||
endpoint = "/v1beta/models/nano-banana-2-1k-official:generateContent"
|
||||
@@ -214,7 +215,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
# 添加文本部分
|
||||
parts.append({"text": prompt})
|
||||
|
||||
image_config = {"imageSize": resolution}
|
||||
image_config = {"imageSize": {"512": "512px"}.get(resolution, resolution)}
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
image_config["aspectRatio"] = aspect_ratio
|
||||
|
||||
@@ -364,7 +365,8 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
async def parse_response_async(
|
||||
self,
|
||||
response: Dict[str, Any],
|
||||
session: Optional[aiohttp.ClientSession] = None
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
image_downloader: Optional[Callable[[str], Awaitable[bytes]]] = None,
|
||||
) -> tuple[List[Image.Image], Dict[str, Any]]:
|
||||
"""
|
||||
异步解析 API 响应，提取生成的图像
|
||||
@@ -432,6 +434,24 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
if session is None:
|
||||
session = self._make_session()
|
||||
close_session = True
|
||||
|
||||
def _tag_image(img: Image.Image, raw_bytes: bytes) -> Image.Image:
|
||||
fmt = (img.format or "").upper()
|
||||
if fmt == "JPG":
|
||||
fmt = "JPEG"
|
||||
if fmt:
|
||||
img.format = fmt
|
||||
setattr(img, "_o1key_original_format", fmt)
|
||||
setattr(img, "_o1key_original_bytes", raw_bytes)
|
||||
return img
|
||||
|
||||
async def _download_url(url: str) -> bytes:
|
||||
if image_downloader is not None:
|
||||
return await image_downloader(url)
|
||||
async with session.get(url) as img_response:
|
||||
if img_response.status != 200:
|
||||
raise RuntimeError(f"图片下载失败 ({img_response.status})")
|
||||
return await img_response.read()
|
||||
|
||||
try:
|
||||
for candidate_idx, candidate in enumerate(candidates):
|
||||
@@ -446,7 +466,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
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:
|
||||
@@ -458,6 +478,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
|
||||
if img_data:
|
||||
img = decode_base64_to_pil(img_data)
|
||||
_tag_image(img, base64.b64decode(img_data))
|
||||
images.append(img)
|
||||
|
||||
# 记录格式信息
|
||||
@@ -475,21 +496,20 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
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_bytes = await _download_url(url)
|
||||
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)
|
||||
img = Image.open(BytesIO(img_bytes))
|
||||
_tag_image(img, 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
|
||||
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 # 静默失败
|
||||
|
||||
@@ -514,46 +534,44 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
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()
|
||||
img_data = await _download_url(url)
|
||||
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))
|
||||
_tag_image(img, 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()
|
||||
img_data = await _download_url(url)
|
||||
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))
|
||||
_tag_image(img, 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 # 静默失败
|
||||
|
||||
|
||||
Reference in New Issue
Block a user