77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
"""Replicate 图像生成服务封装。"""
|
||
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
import httpx
|
||
import replicate
|
||
|
||
GENERATED_DIR = Path(__file__).parent.parent.parent / "generated"
|
||
|
||
|
||
async def _download_image(url: str) -> str:
|
||
"""下载远程图片到本地 generated/ 目录，返回本地 URL 路径。"""
|
||
filename = f"{uuid.uuid4().hex}.png"
|
||
filepath = GENERATED_DIR / filename
|
||
async with httpx.AsyncClient() as client:
|
||
resp = await client.get(url, follow_redirects=True)
|
||
resp.raise_for_status()
|
||
filepath.write_bytes(resp.content)
|
||
return f"/generated/{filename}"
|
||
|
||
|
||
async def generate_images(
|
||
prompt: str,
|
||
num_images: int = 1,
|
||
ref_image_url: str | None = None,
|
||
) -> list[str]:
|
||
"""
|
||
调用 Replicate 生成图片。
|
||
|
||
返回本地可访问的图片 URL 列表。
|
||
"""
|
||
local_urls = []
|
||
|
||
for _ in range(num_images):
|
||
try:
|
||
if ref_image_url and ref_image_url.startswith("/"):
|
||
# 本地路径转为 file URI 不适用于 Replicate，
|
||
# 需要用户上传的图先通过后端 URL 访问
|
||
# MVP 阶段：参考图作为 prompt 的文字补充，不直接传给模型
|
||
# 仅使用 flux-schnell 文生图
|
||
output = await replicate.async_run(
|
||
"black-forest-labs/flux-schnell",
|
||
input={
|
||
"prompt": prompt,
|
||
"num_outputs": 1,
|
||
"aspect_ratio": "1:1",
|
||
"output_format": "png",
|
||
},
|
||
)
|
||
else:
|
||
output = await replicate.async_run(
|
||
"black-forest-labs/flux-schnell",
|
||
input={
|
||
"prompt": prompt,
|
||
"num_outputs": 1,
|
||
"aspect_ratio": "1:1",
|
||
"output_format": "png",
|
||
},
|
||
)
|
||
|
||
# output 是 FileOutput 列表或单个 URL
|
||
if isinstance(output, list):
|
||
for item in output:
|
||
url = str(item)
|
||
local_url = await _download_image(url)
|
||
local_urls.append(local_url)
|
||
else:
|
||
url = str(output)
|
||
local_url = await _download_image(url)
|
||
local_urls.append(local_url)
|
||
|
||
except Exception as e:
|
||
local_urls.append(f"[生成失败: {e}]")
|
||
|
||
return local_urls
|