"""图像生成服务 — 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, )