195 lines
7.2 KiB
Python
195 lines
7.2 KiB
Python
"""
|
||
Gemini 异步生图 Provider
|
||
通过 cf-api.o1key.com 的异步提交+轮询接口调用 Gemini 图像生成模型
|
||
|
||
协议说明：
|
||
- 提交：POST {base}/async{gemini_endpoint}?image_format=url
|
||
- 轮询：GET {base}/async/v1/tasks/{task_id}
|
||
- 结果：可能直接返回 image_url，也可能返回 Gemini 标准 candidates 格式
|
||
"""
|
||
|
||
from io import BytesIO
|
||
from typing import Dict, List, Optional
|
||
from PIL import Image
|
||
|
||
from .base_async_provider import BaseAsyncImageProvider
|
||
from .gemini_client import GeminiAPIClient
|
||
from ..utils.config import get_async_api_base_url, get_api_key_or_raise
|
||
from ..models_config import (
|
||
get_enabled_models,
|
||
get_model_supported_aspect_ratios,
|
||
get_model_supported_resolutions,
|
||
)
|
||
|
||
|
||
class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||
"""
|
||
Gemini 异步生图 Provider
|
||
|
||
委托 GeminiAPIClient 处理：
|
||
- 端点构造（get_endpoint）
|
||
- 请求体构建（build_request_body）
|
||
- 响应解析（parse_response_async）
|
||
"""
|
||
|
||
def __init__(self, api_key: str = None, proxy_url: str = None):
|
||
if api_key is None:
|
||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||
super().__init__(api_key=api_key, proxy_url=proxy_url)
|
||
self._client = GeminiAPIClient(api_key=api_key)
|
||
|
||
# ========================================================================
|
||
# 抽象方法实现
|
||
# ========================================================================
|
||
|
||
@property
|
||
def api_base_url(self) -> str:
|
||
return getattr(self, '_route_base_url', None) or get_async_api_base_url()
|
||
|
||
def get_submit_endpoint(self, model: str, resolution: str) -> str:
|
||
gemini_endpoint = self._client.get_endpoint(
|
||
model=model, resolution=resolution, image_format="url"
|
||
)
|
||
base = gemini_endpoint.split("?")[0]
|
||
async_endpoint = f"/async{base}"
|
||
if "?" in gemini_endpoint:
|
||
async_endpoint += "?" + gemini_endpoint.split("?", 1)[1]
|
||
return async_endpoint
|
||
|
||
def build_submit_body(
|
||
self,
|
||
prompt: str,
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
images: Optional[List[Image.Image]] = None,
|
||
**kwargs
|
||
) -> dict:
|
||
return self._client.build_request_body(
|
||
prompt=prompt,
|
||
images=images,
|
||
aspect_ratio=aspect_ratio,
|
||
resolution=resolution,
|
||
enable_grounding=kwargs.get("enable_grounding", False),
|
||
enable_image_search=kwargs.get("enable_image_search", False),
|
||
image_compression=getattr(self, "image_compression", None),
|
||
thinking_level=kwargs.get("thinking_level"),
|
||
request_log_enabled=False,
|
||
)
|
||
|
||
def extract_task_id(self, response: dict) -> str:
|
||
task_id = response.get("task_id")
|
||
if not task_id:
|
||
raise RuntimeError(f"提交响应中未找到 task_id: {response}")
|
||
return task_id
|
||
|
||
def extract_status(self, response: dict) -> str:
|
||
return response.get("status", "UNKNOWN")
|
||
|
||
async def parse_result(self, result_data: dict, session) -> List[Image.Image]:
|
||
images = result_data.get("images") if isinstance(result_data, dict) else None
|
||
if isinstance(images, list) and images:
|
||
parsed = []
|
||
for item in images:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
image_url = item.get("url") or item.get("image_url")
|
||
if image_url:
|
||
async with session.get(image_url) as img_resp:
|
||
if img_resp.status == 200:
|
||
img_bytes = await img_resp.read()
|
||
parsed.append(Image.open(BytesIO(img_bytes)).convert("RGB"))
|
||
else:
|
||
raise RuntimeError(f"下载图片失败 ({img_resp.status}): {image_url}")
|
||
if parsed:
|
||
return parsed
|
||
|
||
# 异步接口可能直接返回 image_url
|
||
image_url = result_data.get("image_url", "") if isinstance(result_data, dict) else ""
|
||
if image_url:
|
||
async with session.get(image_url) as img_resp:
|
||
if img_resp.status == 200:
|
||
img_bytes = await img_resp.read()
|
||
return [Image.open(BytesIO(img_bytes))]
|
||
raise RuntimeError(f"下载图片失败 ({img_resp.status}): {image_url}")
|
||
|
||
# 否则按 Gemini 标准格式解析
|
||
images_list, _ = await self._client.parse_response_async(result_data, session=session)
|
||
return images_list
|
||
|
||
def get_models(self) -> List[str]:
|
||
return get_enabled_models()
|
||
|
||
def get_model_aspect_ratios(self, model_id: str) -> List[str]:
|
||
return get_model_supported_aspect_ratios(model_id)
|
||
|
||
def get_model_resolutions(self, model_id: str) -> List[str]:
|
||
return get_model_supported_resolutions(model_id)
|
||
|
||
# ========================================================================
|
||
# 可选方法覆盖
|
||
# ========================================================================
|
||
|
||
def get_extra_inputs(self) -> dict:
|
||
"""Gemini 专有：Google Search Grounding"""
|
||
return {
|
||
"联网功能": (["关闭", "打开"], {"default": "关闭"}),
|
||
}
|
||
|
||
def get_extra_kwargs(self, **kwargs) -> dict:
|
||
return {
|
||
"enable_grounding": kwargs.pop("联网功能", "关闭") == "打开",
|
||
}
|
||
|
||
def extract_progress(self, response: dict) -> Optional[float]:
|
||
"""从轮询响应中提取进度（0.0-1.0）"""
|
||
|
||
def _coerce(val) -> Optional[float]:
|
||
if val is None or isinstance(val, bool):
|
||
return None
|
||
if isinstance(val, (int, float)):
|
||
progress = float(val)
|
||
elif isinstance(val, str):
|
||
text = val.strip()
|
||
if not text:
|
||
return None
|
||
has_percent_suffix = text.endswith("%")
|
||
if has_percent_suffix:
|
||
text = text[:-1].strip()
|
||
try:
|
||
progress = float(text)
|
||
except ValueError:
|
||
return None
|
||
if has_percent_suffix:
|
||
progress /= 100.0
|
||
else:
|
||
return None
|
||
if progress > 1.0:
|
||
progress /= 100.0
|
||
return max(0.0, min(progress, 1.0))
|
||
|
||
# 直接字段：progress / percentage
|
||
for field in ("progress", "percentage", "percent"):
|
||
progress = _coerce(response.get(field))
|
||
if progress is not None:
|
||
return progress
|
||
|
||
# 嵌套字段：progressInfo / progress_info
|
||
progress_info = response.get("progressInfo") or response.get("progress_info")
|
||
if isinstance(progress_info, dict):
|
||
for field in ("progress", "percentage", "percent"):
|
||
progress = _coerce(progress_info.get(field))
|
||
if progress is not None:
|
||
return progress
|
||
|
||
return None
|
||
|
||
def query_balance_sync(self) -> Optional[dict]:
|
||
try:
|
||
return self._client.query_balance_sync()
|
||
except Exception:
|
||
return None
|
||
|
||
def format_balance_info(self, balance_data: dict) -> str:
|
||
return self._client.format_balance_info(balance_data)
|