R524809
126 lines
4.7 KiB
Python
126 lines
4.7 KiB
Python
"""图生图 / 图像编辑请求响应模型（多模型）。"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Literal
|
||
|
||
from pydantic import BaseModel, Field, field_validator
|
||
|
||
# 支持的模型白名单
|
||
SUPPORTED_MODELS: frozenset[str] = frozenset(
|
||
{
|
||
"wanx2.1-imageedit",
|
||
"wan2.6-image",
|
||
"qwen-image-edit",
|
||
"qwen-image-edit-plus",
|
||
"qwen-image-edit-plus-2025-10-30",
|
||
}
|
||
)
|
||
|
||
# 万相 2.1 通用图像编辑支持的 function 白名单
|
||
# description_edit ：文本指令编辑（无掩码，即“图生图”）
|
||
# description_edit_with_mask ：局部重绘（需 mask）
|
||
# stylization_local ：局部风格化
|
||
# stylization_all ：整体风格化
|
||
ImageEditFunction = Literal[
|
||
"description_edit",
|
||
"description_edit_with_mask",
|
||
"stylization_local",
|
||
"stylization_all",
|
||
]
|
||
|
||
# 单张 base64 图片（编码后）上限，超出直接拒绝，避免无谓占用内存。
|
||
_MAX_BASE64_LENGTH = 15 * 1024 * 1024
|
||
|
||
|
||
class ImageEditRequest(BaseModel):
|
||
"""图生图请求。图片以 data URL（data:image/...;base64,...）或公网 URL 传入。"""
|
||
|
||
base_image: str = Field(..., description="待编辑原图：data URL 或公网 http(s) URL")
|
||
prompt: str = Field(..., description="编辑指令，例如“把背景换成纯白色摄影棚”")
|
||
model: str = Field(
|
||
default="wanx2.1-imageedit",
|
||
description="模型 code，见 SUPPORTED_MODELS",
|
||
)
|
||
mask_image: str | None = Field(
|
||
default=None, description="掩码图（仅 function=description_edit_with_mask 时使用）"
|
||
)
|
||
function: ImageEditFunction = Field(
|
||
default="description_edit", description="编辑能力类型，见 ImageEditFunction"
|
||
)
|
||
n: int = Field(default=1, ge=1, le=4, description="生成数量")
|
||
size: str = Field(default="", description="输出尺寸，如 1024*1024；留空由模型决定")
|
||
seed: int | None = Field(default=None, ge=0, description="随机种子，可选")
|
||
style: str = Field(default="", description="输出风格，如 auto/photography/anime 等，可选")
|
||
prompt_extend: bool = Field(default=True, description="是否自动扩写提示词（官方示例默认开启）")
|
||
strength: float | None = Field(
|
||
default=None,
|
||
ge=0.0,
|
||
le=1.0,
|
||
description=(
|
||
"修改幅度 0.0~1.0，官方默认 0.5；值越大改动越明显、越偏离原图。"
|
||
"新增文字、大幅修改等场景建议 0.8 左右"
|
||
),
|
||
)
|
||
|
||
@field_validator("prompt")
|
||
@classmethod
|
||
def prompt_not_empty(cls, value: str) -> str:
|
||
text = (value or "").strip()
|
||
if not text:
|
||
raise ValueError("prompt 不能为空")
|
||
return text
|
||
|
||
@field_validator("model")
|
||
@classmethod
|
||
def model_whitelist(cls, value: str) -> str:
|
||
model = (value or "").strip()
|
||
if model not in SUPPORTED_MODELS:
|
||
raise ValueError(f"不支持的模型「{model}」，可选：{sorted(SUPPORTED_MODELS)}")
|
||
return model
|
||
|
||
@field_validator("base_image")
|
||
@classmethod
|
||
def base_image_valid(cls, value: str) -> str:
|
||
image = (value or "").strip()
|
||
if not image:
|
||
raise ValueError("base_image 不能为空")
|
||
if image.startswith("data:image/"):
|
||
if len(image) > _MAX_BASE64_LENGTH:
|
||
raise ValueError("图片过大，请压缩后重试")
|
||
elif not image.startswith(("http://", "https://")):
|
||
raise ValueError("base_image 仅支持 data URL 或公网 http(s) URL")
|
||
return image
|
||
|
||
@field_validator("mask_image")
|
||
@classmethod
|
||
def mask_image_valid(cls, value: str | None) -> str | None:
|
||
if value is None or not value.strip():
|
||
return None
|
||
image = value.strip()
|
||
if image.startswith("data:image/"):
|
||
if len(image) > _MAX_BASE64_LENGTH:
|
||
raise ValueError("mask_image 图片过大")
|
||
elif not image.startswith(("http://", "https://")):
|
||
raise ValueError("mask_image 仅支持 data URL 或公网 http(s) URL")
|
||
return image
|
||
|
||
@field_validator("style")
|
||
@classmethod
|
||
def normalize_style(cls, value: str) -> str:
|
||
return (value or "").strip()
|
||
|
||
|
||
class ImageEditResult(BaseModel):
|
||
url: str = Field(..., description="生成图片 URL（有效期 24 小时）")
|
||
image_base64: str = Field(
|
||
default="",
|
||
description="生成图片的 data URL（服务端代理下载，便于前端直接合成/导出，规避跨域画布污染）",
|
||
)
|
||
|
||
|
||
class ImageEditResponse(BaseModel):
|
||
task_id: str = ""
|
||
results: list[ImageEditResult] = Field(default_factory=list)
|
||
image_count: int = 0
|
||
request_id: str = ""
|