YuShiSheJiShi/backend/app/services/ai_generator.py

"""
AI 生图服务
支持双模型：SiliconFlow Kolors 和 火山引擎 Seedream 5.0 lite
"""
import os
import uuid
import logging
import httpx
from typing import Optional, Tuple

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, seed: Optional[int] = None) -> Tuple[str, Optional[int]]:
    """
    调用 SiliconFlow 生图 API（Kolors 模型）
    
    Args:
        seed: 随机种子，传入相同 seed 可保持多视角图片风格一致
    
    Returns:
        (远程图片 URL, 使用的 seed)
    """
    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": "Kwai-Kolors/Kolors",
        "prompt": prompt,
        "image_size": f"{size}x{size}",
        "batch_size": 1,
        "num_inference_steps": 20,
        "guidance_scale": 7.5,
    }
    if seed is not None:
        payload["seed"] = seed

    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://..."}], "seed": 12345}
    images = data.get("images", [])
    if not images:
        raise ValueError("SiliconFlow 返回空图片列表")
    returned_seed = data.get("seed")
    return images[0]["url"], returned_seed


async def _call_seedream(prompt: str, size: int = 1024, ai_config: dict = None, seed: Optional[int] = None, ref_image_url: Optional[str] = None) -> Tuple[str, Optional[int]]:
    """
    调用火山引擎 Seedream 5.0 lite 生图 API
    
    Args:
        ref_image_url: 参考图 URL，用于多视角一致性（将第一张图作为参考传入后续视角）
    
    Returns:
        (远程图片 URL, seed)
    """
    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-5-0-260128",
        "prompt": prompt,
        "size": "2K",
        "response_format": "url",
        "watermark": False,
    }
    # 传入参考图保持多视角一致性（API 要求数组格式）
    if ref_image_url:
        payload["image"] = [ref_image_url]

    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
        resp = await client.post(url, json=payload, headers=headers)
        if resp.status_code != 200:
            logger.error(f"Seedream API 错误: status={resp.status_code}, body={resp.text[:500]}")
            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"], seed


async def generate_image(prompt: str, model: Optional[str] = None, seed: Optional[int] = None, ref_image_url: Optional[str] = None) -> Tuple[str, Optional[int]]:
    """
    统一生图接口，带重试机制
    
    Args:
        prompt: 提示词
        model: 模型名称 (flux-dev / seedream-5.0)
        seed: 随机种子（SiliconFlow Kolors 支持）
        ref_image_url: 参考图 URL（Seedream 5.0 支持，用于多视角一致性）
        
    Returns:
        (远程图片 URL, 使用的 seed)
    """
    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 in ("seedream-5.0", "seedream-4.5"):
                image_url, returned_seed = await _call_seedream(prompt, size, ai_config, seed, ref_image_url)
            else:
                image_url, returned_seed = await _call_siliconflow(prompt, size, ai_config, seed)
            logger.info(f"AI 生图成功 (model={model}, seed={returned_seed}, attempt={attempt})")
            return image_url, returned_seed
        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}"