221 lines
7.7 KiB
Python
221 lines
7.7 KiB
Python
"""图像生成服务 — Provider 抽象层。"""
|
||
|
||
import base64
|
||
import mimetypes
|
||
import os
|
||
import uuid
|
||
from abc import ABC, abstractmethod
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
from replicate import Client as ReplicateClient
|
||
|
||
from app.config import get_image_model_config, get_image_aspect_ratio, get_image_output_format
|
||
|
||
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
|
||
|
||
|
||
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_url: str | None = None,
|
||
) -> list[str]:
|
||
"""生成图片并返回本地 URL 列表。"""
|
||
...
|
||
|
||
|
||
class ReplicateProvider(ImageProvider):
|
||
"""通过 Replicate API 调用模型。"""
|
||
|
||
async def generate(
|
||
self,
|
||
model_config: dict[str, Any],
|
||
prompt: str,
|
||
num_images: int = 1,
|
||
ref_image_url: 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:
|
||
input_params = self._build_input(
|
||
model_config, default_params, prompt, num_images, ref_image_url
|
||
)
|
||
|
||
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,
|
||
) -> dict[str, Any]:
|
||
"""
|
||
根据模型配置构建 Replicate input 参数。
|
||
|
||
通过 model_config 中的标志字段自动适配不同模型：
|
||
- supports_ref_image / ref_image_param: 参考图注入
|
||
- num_images_param: 各模型的批量生成参数名（如 number_of_images / num_outputs）
|
||
"""
|
||
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")
|
||
|
||
# ── 支持参考图的模型（IP-Adapter 系列）──
|
||
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
|
||
|
||
# ── Flux 系列（纯文生图）──
|
||
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:
|
||
# SDXL 类模型
|
||
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"]
|
||
|
||
return params
|
||
|
||
|
||
# ─── Provider 注册 ─────────────────────────────────────
|
||
|
||
_PROVIDERS: dict[str, ImageProvider] = {
|
||
"replicate": ReplicateProvider(),
|
||
}
|
||
|
||
|
||
# ─── 公开 API ──────────────────────────────────────────
|
||
|
||
class GenerateResult:
|
||
"""图片生成结果，包含生成的 URL 列表和实际使用的模型信息。"""
|
||
|
||
def __init__(self, urls: list[str], model_name: str, model_id: str):
|
||
self.urls = urls
|
||
self.model_name = model_name
|
||
self.model_id = model_id
|
||
|
||
|
||
async def generate_images(
|
||
prompt: str,
|
||
num_images: int = 1,
|
||
ref_image_url: str | None = None,
|
||
model_id: str | None = None,
|
||
) -> GenerateResult:
|
||
"""
|
||
统一入口：根据 model_id 查注册表，分发到对应 provider。
|
||
|
||
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)
|
||
|
||
urls = await provider.generate(config, prompt, num_images, ref_image_url)
|
||
return GenerateResult(urls, model_name, resolved_id)
|