shiyong.qin/EPEEAIKit
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

加了一堆模型和一堆功能

Browse Source
This commit is contained in:
Shin@HOME
2026-04-16 00:36:21 +08:00
parent 47c0863bab
commit 40de992516
34 changed files with 3111 additions and 718 deletions
Download Patch File Download Diff File Expand all files Collapse all files

383
art-agent/backend/app/services/image_gen.py
View File

@@ -1,5 +1,6 @@
"""图像生成服务 — Provider 抽象层。"""
import asyncio
import base64
import mimetypes
import os
@@ -9,9 +10,11 @@ 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_model_config, get_image_aspect_ratio, get_image_output_format
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"
@@ -57,6 +60,43 @@ def to_data_uri(image_path: str) -> str:
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"
@@ -80,9 +120,10 @@ class ImageProvider(ABC):
model_config: dict[str, Any],
prompt: str,
num_images: int = 1,
ref_image_url: str | None = None,
ref_image_urls: list[str] | None = None,
negative_prompt: str | None = None,
) -> list[str]:
"""生成图片并返回本地 URL 列表。"""
"""生成图片并返回本地 URL 列表。negative_prompt 仅部分后端使用（如 Replicate SDXL）。"""
...
@@ -94,7 +135,8 @@ class ReplicateProvider(ImageProvider):
model_config: dict[str, Any],
prompt: str,
num_images: int = 1,
ref_image_url: str | None = None,
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", {})
@@ -102,8 +144,15 @@ class ReplicateProvider(ImageProvider):
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, ref_image_url
model_config,
default_params,
prompt,
num_images,
single_ref,
negative_prompt=negative_prompt,
)
output = await _replicate_client.async_run(
@@ -131,20 +180,18 @@ class ReplicateProvider(ImageProvider):
prompt: str,
num_images: int = 1,
ref_image_url: str | None = None,
negative_prompt: str | None = None,
) -> dict[str, Any]:
"""
根据模型配置构建 Replicate input 参数。
通过 model_config 中的标志字段自动适配不同模型：
- supports_ref_image / ref_image_param: 参考图注入
- num_images_param: 各模型的批量生成参数名（如 number_of_images / num_outputs）
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")
# ── 支持参考图的模型（IP-Adapter 系列）──
if supports_ref:
params: dict[str, Any] = {"prompt": prompt}
for k, v in default_params.items():
@@ -154,7 +201,6 @@ class ReplicateProvider(ImageProvider):
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}
@@ -166,21 +212,281 @@ class ReplicateProvider(ImageProvider):
"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"]
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(),
}
@@ -189,21 +495,32 @@ _PROVIDERS: dict[str, ImageProvider] = {
class GenerateResult:
"""图片生成结果，包含生成的 URL 列表和实际使用的模型信息。"""
def __init__(self, urls: list[str], model_name: str, model_id: str):
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_url: str | None = None,
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)
@@ -216,5 +533,41 @@ async def generate_images(
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)
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,
)