574 lines
21 KiB
Python
574 lines
21 KiB
Python
"""图像生成服务 — Provider 抽象层。"""
|
||
|
||
import asyncio
|
||
import base64
|
||
import mimetypes
|
||
import os
|
||
import uuid
|
||
from abc import ABC, abstractmethod
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
from openai import AsyncOpenAI
|
||
from replicate import Client as ReplicateClient
|
||
|
||
from app.config import get_image_aspect_ratio, get_image_model_config, get_image_output_format
|
||
from app.services.image_prompt_strategy import prepare_image_prompt_for_model
|
||
|
||
BACKEND_ROOT = Path(__file__).parent.parent.parent
|
||
GENERATED_DIR = BACKEND_ROOT / "generated"
|
||
|
||
|
||
def _make_replicate_client() -> ReplicateClient:
|
||
"""
|
||
创建 Replicate 客户端。
|
||
Replicate SDK 内部创建 httpx transport 时不读取代理环境变量，
|
||
需要我们手动把代理配置注入到 transport 中。
|
||
"""
|
||
timeout = httpx.Timeout(5.0, read=300.0, write=30.0, connect=30.0, pool=10.0)
|
||
proxy_url = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY")
|
||
|
||
transport_kwargs: dict[str, Any] = {}
|
||
if proxy_url:
|
||
transport_kwargs["proxy"] = proxy_url
|
||
|
||
our_transport = httpx.AsyncHTTPTransport(**transport_kwargs)
|
||
|
||
client = ReplicateClient(
|
||
timeout=timeout,
|
||
transport=our_transport,
|
||
)
|
||
|
||
return client
|
||
|
||
|
||
_replicate_client = _make_replicate_client()
|
||
|
||
|
||
# ─── 公共工具 ──────────────────────────────────────────
|
||
|
||
def to_data_uri(image_path: str) -> str:
|
||
"""将本地路径或已有 URL 转为 Replicate 可接受的格式（data URI 或原始 URL）。"""
|
||
if image_path.startswith("data:") or image_path.startswith("http"):
|
||
return image_path
|
||
local = BACKEND_ROOT / image_path.lstrip("/")
|
||
if local.exists():
|
||
mime = mimetypes.guess_type(str(local))[0] or "image/png"
|
||
b64 = base64.b64encode(local.read_bytes()).decode()
|
||
return f"data:{mime};base64,{b64}"
|
||
return image_path
|
||
|
||
|
||
def _load_image_bytes(image_path: str) -> bytes:
|
||
"""将本地路径或 data URI 转为原始字节，供 OpenAI images.edit 使用。"""
|
||
if image_path.startswith("data:"):
|
||
# data:image/png;base64,xxxx
|
||
_, b64_part = image_path.split(",", 1)
|
||
return base64.b64decode(b64_part)
|
||
if image_path.startswith("http"):
|
||
raise ValueError("_load_image_bytes 不支持远程 URL，请先下载到本地")
|
||
local = BACKEND_ROOT / image_path.lstrip("/")
|
||
if local.exists():
|
||
return local.read_bytes()
|
||
raise FileNotFoundError(f"参考图文件未找到: {local}")
|
||
|
||
|
||
def _resolve_image_base64(image_path: str) -> tuple[str, str]:
|
||
"""将图片路径/data URI 解析为 (mime_type, base64_string)。
|
||
|
||
统一处理三种输入形式：
|
||
- data URI (data:image/png;base64,xxxx) → 直接提取 mime 和 base64
|
||
- 本地路径 (/uploads/xxx.png) → 读取文件并编码
|
||
- 其他 → 尝试作为本地路径处理
|
||
"""
|
||
if image_path.startswith("data:"):
|
||
header, b64_part = image_path.split(",", 1)
|
||
# header 格式: data:image/png;base64
|
||
mime = header.split(";")[0].replace("data:", "")
|
||
return mime, b64_part
|
||
|
||
local = BACKEND_ROOT / image_path.lstrip("/")
|
||
if local.exists():
|
||
mime = mimetypes.guess_type(str(local))[0] or "image/png"
|
||
b64 = base64.b64encode(local.read_bytes()).decode()
|
||
return mime, b64
|
||
|
||
raise FileNotFoundError(f"参考图文件未找到: {local}")
|
||
|
||
|
||
async def _download_image(url: str) -> str:
|
||
"""下载远程图片到本地 generated/ 目录，返回本地 URL 路径。"""
|
||
filename = f"{uuid.uuid4().hex}.png"
|
||
filepath = GENERATED_DIR / filename
|
||
proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY")
|
||
async with httpx.AsyncClient(proxy=proxy, timeout=httpx.Timeout(60.0)) as client:
|
||
resp = await client.get(url, follow_redirects=True)
|
||
resp.raise_for_status()
|
||
filepath.write_bytes(resp.content)
|
||
return f"/generated/{filename}"
|
||
|
||
|
||
# ─── Provider 抽象基类 ─────────────────────────────────
|
||
|
||
class ImageProvider(ABC):
|
||
"""所有图像生成 provider 的基类。"""
|
||
|
||
@abstractmethod
|
||
async def generate(
|
||
self,
|
||
model_config: dict[str, Any],
|
||
prompt: str,
|
||
num_images: int = 1,
|
||
ref_image_urls: list[str] | None = None,
|
||
negative_prompt: str | None = None,
|
||
) -> list[str]:
|
||
"""生成图片并返回本地 URL 列表。negative_prompt 仅部分后端使用（如 Replicate SDXL）。"""
|
||
...
|
||
|
||
|
||
class ReplicateProvider(ImageProvider):
|
||
"""通过 Replicate API 调用模型。"""
|
||
|
||
async def generate(
|
||
self,
|
||
model_config: dict[str, Any],
|
||
prompt: str,
|
||
num_images: int = 1,
|
||
ref_image_urls: list[str] | None = None,
|
||
negative_prompt: str | None = None,
|
||
) -> list[str]:
|
||
replicate_model_id = model_config["model_id"]
|
||
default_params: dict[str, Any] = model_config.get("default_params", {})
|
||
|
||
local_urls: list[str] = []
|
||
|
||
try:
|
||
# Replicate 模型（InstantStyle/Kolors）只支持单张参考图，取第一张
|
||
single_ref = ref_image_urls[0] if ref_image_urls else None
|
||
input_params = self._build_input(
|
||
model_config,
|
||
default_params,
|
||
prompt,
|
||
num_images,
|
||
single_ref,
|
||
negative_prompt=negative_prompt,
|
||
)
|
||
|
||
output = await _replicate_client.async_run(
|
||
replicate_model_id, input=input_params, wait=False
|
||
)
|
||
|
||
items = output if isinstance(output, list) else [output]
|
||
for item in items:
|
||
url = str(item)
|
||
if url.startswith("https://") or url.startswith("http://") or url.startswith("data:"):
|
||
local_urls.append(await _download_image(url))
|
||
else:
|
||
local_urls.append(f"[生成失败: 模型返回非图片内容: {url[:200]}]")
|
||
|
||
except Exception as e:
|
||
detail = str(e) or f"{type(e).__name__}: {repr(e)}"
|
||
local_urls.append(f"[生成失败: {detail}]")
|
||
|
||
return local_urls
|
||
|
||
@staticmethod
|
||
def _build_input(
|
||
model_config: dict[str, Any],
|
||
default_params: dict[str, Any],
|
||
prompt: str,
|
||
num_images: int = 1,
|
||
ref_image_url: str | None = None,
|
||
negative_prompt: str | None = None,
|
||
) -> dict[str, Any]:
|
||
"""
|
||
根据模型配置构建 Replicate input 参数。
|
||
|
||
Replicate 上的 IP-Adapter 模型只支持单张参考图。
|
||
"""
|
||
model_id = model_config["model_id"]
|
||
supports_ref = model_config.get("supports_ref_image", False)
|
||
ref_param_name = model_config.get("ref_image_param", "image")
|
||
num_images_param = model_config.get("num_images_param", "num_outputs")
|
||
|
||
if supports_ref:
|
||
params: dict[str, Any] = {"prompt": prompt}
|
||
for k, v in default_params.items():
|
||
params[k] = v
|
||
params[num_images_param] = num_images
|
||
if ref_image_url:
|
||
params[ref_param_name] = to_data_uri(ref_image_url)
|
||
return params
|
||
|
||
is_flux = "flux" in model_id.lower()
|
||
params = {"prompt": prompt, num_images_param: num_images}
|
||
|
||
if is_flux:
|
||
params["aspect_ratio"] = default_params.get(
|
||
"aspect_ratio", get_image_aspect_ratio()
|
||
)
|
||
params["output_format"] = default_params.get(
|
||
"output_format", get_image_output_format()
|
||
)
|
||
else:
|
||
params["width"] = default_params.get("width", 1024)
|
||
params["height"] = default_params.get("height", 1024)
|
||
if "num_inference_steps" in default_params:
|
||
params["num_inference_steps"] = default_params["num_inference_steps"]
|
||
if "guidance_scale" in default_params:
|
||
params["guidance_scale"] = default_params["guidance_scale"]
|
||
if negative_prompt is not None:
|
||
params["negative_prompt"] = negative_prompt
|
||
|
||
return params
|
||
|
||
|
||
class OpenAIImageProvider(ImageProvider):
|
||
"""通过 OpenAI 兼容 API（向量引擎中转）调用 GPT Image 系列模型。
|
||
|
||
有参考图时使用 images.edit（支持最多 16 张参考图），
|
||
无参考图时使用 images.generate（纯文生图）。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self._client: AsyncOpenAI | None = None
|
||
|
||
def _get_client(self) -> AsyncOpenAI:
|
||
if self._client is None:
|
||
self._client = AsyncOpenAI(
|
||
api_key=os.getenv("VECTORENGINE_API_KEY"),
|
||
base_url=os.getenv("VECTORENGINE_BASE_URL", "https://api.vectorengine.ai/v1"),
|
||
timeout=httpx.Timeout(5.0, read=180.0, write=60.0, connect=30.0),
|
||
)
|
||
return self._client
|
||
|
||
async def generate(
|
||
self,
|
||
model_config: dict[str, Any],
|
||
prompt: str,
|
||
num_images: int = 1,
|
||
ref_image_urls: list[str] | None = None,
|
||
negative_prompt: str | None = None,
|
||
) -> list[str]:
|
||
_ = negative_prompt
|
||
client = self._get_client()
|
||
default_params = model_config.get("default_params", {})
|
||
model_id = model_config["model_id"]
|
||
|
||
local_urls: list[str] = []
|
||
try:
|
||
if ref_image_urls:
|
||
resp = await self._edit_with_refs(
|
||
client, model_id, prompt, ref_image_urls, num_images, default_params
|
||
)
|
||
else:
|
||
resp = await client.images.generate(
|
||
model=model_id,
|
||
prompt=prompt,
|
||
n=num_images,
|
||
size=default_params.get("size", "1024x1024"),
|
||
quality=default_params.get("quality", "high"),
|
||
)
|
||
|
||
for img_data in resp.data:
|
||
if img_data.url:
|
||
local_urls.append(await _download_image(img_data.url))
|
||
elif img_data.b64_json:
|
||
filename = f"{uuid.uuid4().hex}.png"
|
||
filepath = GENERATED_DIR / filename
|
||
filepath.write_bytes(base64.b64decode(img_data.b64_json))
|
||
local_urls.append(f"/generated/{filename}")
|
||
else:
|
||
local_urls.append("[生成失败: 模型未返回图片数据]")
|
||
|
||
except Exception as e:
|
||
detail = str(e) or f"{type(e).__name__}: {repr(e)}"
|
||
local_urls.append(f"[生成失败: {detail}]")
|
||
|
||
return local_urls
|
||
|
||
@staticmethod
|
||
async def _edit_with_refs(
|
||
client: AsyncOpenAI,
|
||
model_id: str,
|
||
prompt: str,
|
||
ref_image_urls: list[str],
|
||
num_images: int,
|
||
default_params: dict[str, Any],
|
||
):
|
||
"""使用 images.edit 端点传入参考图（GPT Image 系列最多 16 张）。"""
|
||
image_files: list[Any] = []
|
||
for url in ref_image_urls:
|
||
image_files.append(_load_image_bytes(url))
|
||
|
||
image_arg: Any = image_files[0] if len(image_files) == 1 else image_files
|
||
|
||
return await client.images.edit(
|
||
model=model_id,
|
||
image=image_arg,
|
||
prompt=prompt,
|
||
n=num_images,
|
||
size=default_params.get("size", "1024x1024"),
|
||
)
|
||
|
||
|
||
class GeminiNativeImageProvider(ImageProvider):
|
||
"""通过向量引擎中转调用 Gemini 原生 generateContent 接口。
|
||
|
||
Gemini 原生接口支持文字 + 图片混合输入（最多 14 张参考图），
|
||
在一次 generateContent 调用中同时理解参考图并生成新图片。
|
||
这是 OpenAI 兼容的 images/generate 和 images/edit 都无法覆盖的能力。
|
||
|
||
API 格式：
|
||
POST /v1beta/models/{model}:generateContent?key={API_KEY}
|
||
Body: { contents: [{ parts: [...] }], generationConfig: { responseModalities: ["TEXT","IMAGE"] } }
|
||
Response: candidates[0].content.parts[] → text 或 inline_data (base64)
|
||
|
||
多图 + 生图耗时较长，上游或代理可能提前断开（httpx: Server disconnected without sending a response）。
|
||
使用较长超时 + 对可恢复网络错误自动重试。
|
||
"""
|
||
|
||
_RETRYABLE: tuple[type[BaseException], ...] = (
|
||
httpx.RemoteProtocolError,
|
||
httpx.ConnectError,
|
||
httpx.ReadTimeout,
|
||
httpx.WriteTimeout,
|
||
httpx.ConnectTimeout,
|
||
httpx.PoolTimeout,
|
||
)
|
||
|
||
@classmethod
|
||
def _gemini_timeout(cls) -> httpx.Timeout:
|
||
"""可通过环境变量调大，多参考图时请求体大、响应慢。"""
|
||
read_s = float(os.getenv("VECTORENGINE_GEMINI_READ_TIMEOUT", "600"))
|
||
write_s = float(os.getenv("VECTORENGINE_GEMINI_WRITE_TIMEOUT", "180"))
|
||
connect_s = float(os.getenv("VECTORENGINE_GEMINI_CONNECT_TIMEOUT", "60"))
|
||
pool_s = float(os.getenv("VECTORENGINE_GEMINI_POOL_TIMEOUT", "60"))
|
||
return httpx.Timeout(
|
||
connect=connect_s,
|
||
read=read_s,
|
||
write=write_s,
|
||
pool=pool_s,
|
||
)
|
||
|
||
async def generate(
|
||
self,
|
||
model_config: dict[str, Any],
|
||
prompt: str,
|
||
num_images: int = 1,
|
||
ref_image_urls: list[str] | None = None,
|
||
negative_prompt: str | None = None,
|
||
) -> list[str]:
|
||
_ = negative_prompt
|
||
api_key = os.getenv("VECTORENGINE_API_KEY", "")
|
||
base_url = os.getenv("VECTORENGINE_BASE_URL", "https://api.vectorengine.ai/v1")
|
||
# 从 /v1 回退到根 URL，拼接 /v1beta/models/... 端点
|
||
api_root = base_url.rstrip("/").removesuffix("/v1")
|
||
model_id = model_config["model_id"]
|
||
|
||
url = f"{api_root}/v1beta/models/{model_id}:generateContent?key={api_key}"
|
||
|
||
parts = self._build_parts(prompt, ref_image_urls)
|
||
body = {
|
||
"contents": [{"parts": parts}],
|
||
"generationConfig": {
|
||
"responseModalities": ["TEXT", "IMAGE"],
|
||
},
|
||
}
|
||
|
||
local_urls: list[str] = []
|
||
proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY")
|
||
max_retries = max(1, int(os.getenv("VECTORENGINE_GEMINI_MAX_RETRIES", "3")))
|
||
timeout = self._gemini_timeout()
|
||
|
||
try:
|
||
data: dict[str, Any] | None = None
|
||
for attempt in range(max_retries):
|
||
try:
|
||
async with httpx.AsyncClient(
|
||
proxy=proxy,
|
||
timeout=timeout,
|
||
limits=httpx.Limits(max_keepalive_connections=5, max_connections=10),
|
||
) as client:
|
||
resp = await client.post(
|
||
url,
|
||
json=body,
|
||
headers={"Content-Type": "application/json"},
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
break
|
||
except self._RETRYABLE as e:
|
||
if attempt < max_retries - 1:
|
||
await asyncio.sleep(2 ** attempt)
|
||
continue
|
||
raise
|
||
|
||
if data is None:
|
||
local_urls.append("[生成失败: 未收到上游响应]")
|
||
return local_urls
|
||
|
||
local_urls = self._extract_images(data)
|
||
|
||
if not local_urls:
|
||
# Gemini 可能只返回了文字（拒绝生图或纯文字回复）
|
||
text_parts = self._extract_text(data)
|
||
hint = text_parts[:200] if text_parts else "模型未返回图片"
|
||
local_urls.append(f"[生成失败: {hint}]")
|
||
|
||
except httpx.HTTPStatusError as e:
|
||
detail = e.response.text[:500] if e.response else str(e)
|
||
local_urls.append(f"[生成失败: HTTP {e.response.status_code} - {detail}]")
|
||
except self._RETRYABLE as e:
|
||
hint = (
|
||
"连接被上游或代理提前关闭，常见于多图参考或生图较慢。"
|
||
"可稍后重试，或在 .env 中增大 VECTORENGINE_GEMINI_READ_TIMEOUT / 检查代理稳定性。"
|
||
)
|
||
detail = str(e) or f"{type(e).__name__}: {repr(e)}"
|
||
local_urls.append(f"[生成失败: {detail}。{hint}]")
|
||
except Exception as e:
|
||
detail = str(e) or f"{type(e).__name__}: {repr(e)}"
|
||
local_urls.append(f"[生成失败: {detail}]")
|
||
|
||
return local_urls
|
||
|
||
@staticmethod
|
||
def _build_parts(
|
||
prompt: str, ref_image_urls: list[str] | None
|
||
) -> list[dict[str, Any]]:
|
||
"""构建 Gemini generateContent 的 parts 数组：文字 + 内联图片。"""
|
||
parts: list[dict[str, Any]] = [{"text": prompt}]
|
||
if ref_image_urls:
|
||
for url in ref_image_urls:
|
||
mime, b64 = _resolve_image_base64(url)
|
||
parts.append({
|
||
"inline_data": {
|
||
"mime_type": mime,
|
||
"data": b64,
|
||
}
|
||
})
|
||
return parts
|
||
|
||
@staticmethod
|
||
def _extract_images(response_data: dict) -> list[str]:
|
||
"""从 Gemini 响应中提取所有图片并保存到本地。"""
|
||
local_urls: list[str] = []
|
||
candidates = response_data.get("candidates", [])
|
||
for candidate in candidates:
|
||
parts = candidate.get("content", {}).get("parts", [])
|
||
for part in parts:
|
||
inline = part.get("inlineData") or part.get("inline_data")
|
||
if inline and inline.get("data"):
|
||
mime = inline.get("mimeType") or inline.get("mime_type", "image/png")
|
||
ext = ".png" if "png" in mime else ".jpg" if "jpeg" in mime or "jpg" in mime else ".webp" if "webp" in mime else ".png"
|
||
filename = f"{uuid.uuid4().hex}{ext}"
|
||
filepath = GENERATED_DIR / filename
|
||
filepath.write_bytes(base64.b64decode(inline["data"]))
|
||
local_urls.append(f"/generated/{filename}")
|
||
return local_urls
|
||
|
||
@staticmethod
|
||
def _extract_text(response_data: dict) -> str:
|
||
"""从 Gemini 响应中提取文字内容（用于调试或错误提示）。"""
|
||
texts: list[str] = []
|
||
candidates = response_data.get("candidates", [])
|
||
for candidate in candidates:
|
||
parts = candidate.get("content", {}).get("parts", [])
|
||
for part in parts:
|
||
if part.get("text"):
|
||
texts.append(part["text"])
|
||
return "\n".join(texts)
|
||
|
||
|
||
# ─── Provider 注册 ─────────────────────────────────────
|
||
|
||
_PROVIDERS: dict[str, ImageProvider] = {
|
||
"replicate": ReplicateProvider(),
|
||
"openai": OpenAIImageProvider(),
|
||
"gemini_native": GeminiNativeImageProvider(),
|
||
}
|
||
|
||
|
||
# ─── 公开 API ──────────────────────────────────────────
|
||
|
||
class GenerateResult:
|
||
"""图片生成结果，包含生成的 URL 列表和实际使用的模型信息。"""
|
||
|
||
def __init__(
|
||
self,
|
||
urls: list[str],
|
||
model_name: str,
|
||
model_id: str,
|
||
*,
|
||
effective_prompt: str | None = None,
|
||
negative_prompt: str | None = None,
|
||
):
|
||
self.urls = urls
|
||
self.model_name = model_name
|
||
self.model_id = model_id
|
||
self.effective_prompt = effective_prompt
|
||
self.negative_prompt = negative_prompt
|
||
|
||
|
||
async def generate_images(
|
||
prompt: str,
|
||
num_images: int = 1,
|
||
ref_image_urls: list[str] | None = None,
|
||
model_id: str | None = None,
|
||
) -> GenerateResult:
|
||
"""
|
||
统一入口：根据 model_id 查注册表，分发到对应 provider。
|
||
|
||
ref_image_urls: 参考图路径列表（可为 None 或空列表）。
|
||
model_id 为空时使用 .env 中配置的默认模型。
|
||
"""
|
||
config = get_image_model_config(model_id)
|
||
provider_name = config.get("provider", "replicate")
|
||
provider = _PROVIDERS.get(provider_name)
|
||
|
||
model_name = config.get("name", "未知模型")
|
||
resolved_id = config.get("id", model_id or "unknown")
|
||
|
||
if not provider:
|
||
return GenerateResult([f"[未知 provider: {provider_name}]"], model_name, resolved_id)
|
||
|
||
effective_refs = ref_image_urls or []
|
||
|
||
# Replicate IP-Adapter 系列模型必须有参考图才能工作
|
||
if config.get("supports_ref_image") and not effective_refs and provider_name == "replicate":
|
||
return GenerateResult(
|
||
[f"[生成失败: {model_name} 是风格迁移模型，需要上传参考图才能使用]"],
|
||
model_name,
|
||
resolved_id,
|
||
)
|
||
|
||
prepared = prepare_image_prompt_for_model(
|
||
config,
|
||
prompt,
|
||
has_reference_images=bool(effective_refs),
|
||
)
|
||
if not prepared.prompt.strip():
|
||
return GenerateResult(
|
||
[f"[生成失败: 经模型策略处理后的 prompt 为空]"],
|
||
model_name,
|
||
resolved_id,
|
||
effective_prompt=prepared.prompt,
|
||
negative_prompt=prepared.negative_prompt,
|
||
)
|
||
|
||
urls = await provider.generate(
|
||
config,
|
||
prepared.prompt,
|
||
num_images,
|
||
effective_refs or None,
|
||
negative_prompt=prepared.negative_prompt,
|
||
)
|
||
return GenerateResult(
|
||
urls,
|
||
model_name,
|
||
resolved_id,
|
||
effective_prompt=prepared.prompt,
|
||
negative_prompt=prepared.negative_prompt,
|
||
)
|