YuShiSheJiShi/backend/app/services/ai_generator.py

"""
AI 生图服务
支持双模型：SiliconFlow FLUX.1 [dev] 和 火山引擎 Seedream 4.5
"""
import os
import uuid
import logging
import httpx
from typing import Optional

from ..config import settings
from .config_service import get_ai_config

logger = logging.getLogger(__name__)

# 超时设置（秒）
REQUEST_TIMEOUT = 90
# 最大重试次数
MAX_RETRIES = 3


async def _call_siliconflow(prompt: str, size: int = 1024, ai_config: dict = None) -> str:
    """
    调用 SiliconFlow FLUX.1 [dev] 生图 API
    
    Returns:
        远程图片 URL
    """
    cfg = ai_config or get_ai_config()
    url = f"{cfg['SILICONFLOW_BASE_URL']}/images/generations"
    headers = {
        "Authorization": f"Bearer {cfg['SILICONFLOW_API_KEY']}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": "black-forest-labs/FLUX.1-dev",
        "prompt": prompt,
        "image_size": f"{size}x{size}",
        "num_inference_steps": 20,
    }

    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
        resp = await client.post(url, json=payload, headers=headers)
        resp.raise_for_status()
        data = resp.json()

    # SiliconFlow 响应格式: {"images": [{"url": "https://..."}]}
    images = data.get("images", [])
    if not images:
        raise ValueError("SiliconFlow 返回空图片列表")
    return images[0]["url"]


async def _call_seedream(prompt: str, size: int = 1024, ai_config: dict = None) -> str:
    """
    调用火山引擎 Seedream 4.5 生图 API
    
    Returns:
        远程图片 URL
    """
    cfg = ai_config or get_ai_config()
    url = f"{cfg['VOLCENGINE_BASE_URL']}/images/generations"
    headers = {
        "Authorization": f"Bearer {cfg['VOLCENGINE_API_KEY']}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": "doubao-seedream-4.5-t2i-250528",
        "prompt": prompt,
        "size": f"{size}x{size}",
        "response_format": "url",
    }

    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
        resp = await client.post(url, json=payload, headers=headers)
        resp.raise_for_status()
        data = resp.json()

    # Seedream 响应格式: {"data": [{"url": "https://..."}]}
    items = data.get("data", [])
    if not items:
        raise ValueError("Seedream 返回空图片列表")
    return items[0]["url"]


async def generate_image(prompt: str, model: Optional[str] = None) -> str:
    """
    统一生图接口，带重试机制
    
    Args:
        prompt: 英文提示词
        model: 模型名称 (flux-dev / seedream-4.5)，为空则使用配置默认值
        
    Returns:
        远程图片 URL
        
    Raises:
        Exception: 所有重试失败后抛出
    """
    ai_config = get_ai_config()
    model = model or ai_config.get("AI_IMAGE_MODEL", "flux-dev")
    size = ai_config.get("AI_IMAGE_SIZE", 1024)

    last_error: Optional[Exception] = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            if model == "seedream-4.5":
                image_url = await _call_seedream(prompt, size, ai_config)
            else:
                image_url = await _call_siliconflow(prompt, size, ai_config)
            logger.info(f"AI 生图成功 (model={model}, attempt={attempt})")
            return image_url
        except Exception as e:
            last_error = e
            logger.warning(f"AI 生图失败 (model={model}, attempt={attempt}/{MAX_RETRIES}): {e}")
            if attempt < MAX_RETRIES:
                import asyncio
                await asyncio.sleep(2 * attempt)  # 指数退避

    raise RuntimeError(f"AI 生图在 {MAX_RETRIES} 次重试后仍然失败: {last_error}")


async def download_and_save(image_url: str, save_path: str) -> str:
    """
    下载远程图片并保存到本地
    
    Args:
        image_url: 远程图片 URL
        save_path: 本地保存路径（如 uploads/designs/1001_效果图.png）
        
    Returns:
        本地文件相对路径（以 / 开头，如 /uploads/designs/1001_效果图.png）
    """
    # 确保目录存在
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client:
        resp = await client.get(image_url)
        resp.raise_for_status()
        with open(save_path, "wb") as f:
            f.write(resp.content)

    logger.info(f"图片已下载保存: {save_path}")
    return f"/{save_path}"