大活

art-agent/backend/app/services/image_gen.py

@@ -0,0 +1,76 @@
+"""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