web-page-backend/music-lab/app/pipeline/cover.py

"""AI 커버 아트 생성 — DALL·E 3 / gpt-image-1 + 그라데이션 폴백."""
import base64
import logging
import os
from io import BytesIO

import httpx
from PIL import Image

from . import storage
from .gradient import make_gradient_with_title

logger = logging.getLogger("music-lab.cover")

DALLE_TIMEOUT_S = 90
PEXELS_IMG_TIMEOUT_S = 30


def _get_api_key() -> str:
    return os.getenv("OPENAI_API_KEY", "")


def _get_model() -> str:
    return os.getenv("OPENAI_IMAGE_MODEL", "gpt-image-1")


def _get_pexels_key() -> str:
    return os.getenv("PEXELS_API_KEY", "")


async def _generate_with_pexels(genre: str, mood: str, track_title: str,
                                  out_path: str, keyword_override: str = "") -> bool:
    """Pexels 이미지 검색·다운로드. 성공 시 True. API key 없거나 0 결과면 False."""
    api_key = _get_pexels_key()
    if not api_key:
        return False
    keyword = keyword_override or f"{genre} aesthetic background"
    try:
        async with httpx.AsyncClient(timeout=PEXELS_IMG_TIMEOUT_S) as client:
            resp = await client.get(
                "https://api.pexels.com/v1/search",
                headers={"Authorization": api_key},
                params={"query": keyword, "per_page": 5, "orientation": "landscape"},
            )
            resp.raise_for_status()
            data = resp.json()
            photos = data.get("photos", [])
            if not photos:
                return False
            img_url = photos[0]["src"].get("large2x") or photos[0]["src"].get("original")
            img_resp = await client.get(img_url)
            img_resp.raise_for_status()
            with Image.open(BytesIO(img_resp.content)) as src:
                img = src.convert("RGB")
            img.save(out_path, "JPEG", quality=92)
            return True
    except (httpx.HTTPError, httpx.TimeoutException, KeyError, ValueError, OSError) as e:
        logger.warning("Pexels 이미지 검색 실패: %s", e)
        return False


async def generate(*, pipeline_id: int, genre: str, prompt_template: str,
                   mood: str = "", track_title: str = "", feedback: str = "",
                   image_source: str = "ai",
                   background_keyword: str = "") -> dict:
    """커버 아트 생성. 성공 시 jpg 저장 + URL 반환. 실패 시 그라데이션 폴백.

    image_source: 'ai' (DALL·E 기본) | 'pexels' (스톡 사진).
    반환: {"url": str, "used_fallback": bool, "error": str | None}
    """
    out_path = os.path.join(storage.pipeline_dir(pipeline_id), "cover.jpg")

    if image_source == "pexels":
        ok = await _generate_with_pexels(genre, mood, track_title, out_path, background_keyword)
        if ok:
            return {
                "url": storage.media_url(pipeline_id, "cover.jpg"),
                "used_fallback": False,
                "error": None,
            }
        # Pexels 실패 → 그라데이션 폴백
        make_gradient_with_title(genre, track_title, out_path)
        return {
            "url": storage.media_url(pipeline_id, "cover.jpg"),
            "used_fallback": True,
            "error": "Pexels 검색 실패 또는 API 키 없음",
        }

    used_fallback = False
    error = None

    api_key = _get_api_key()
    model = _get_model()
    if api_key:
        try:
            await _generate_with_dalle(prompt_template, mood, feedback, out_path,
                                       api_key=api_key, model=model,
                                       background_keyword=background_keyword)
        except (httpx.HTTPError, httpx.TimeoutException, KeyError, ValueError, OSError) as e:
            logger.warning("DALL·E 실패 — 폴백: %s", e)
            error = str(e)
            used_fallback = True
            make_gradient_with_title(genre, track_title, out_path)
    else:
        used_fallback = True
        error = "OPENAI_API_KEY 미설정"
        make_gradient_with_title(genre, track_title, out_path)

    return {
        "url": storage.media_url(pipeline_id, "cover.jpg"),
        "used_fallback": used_fallback,
        "error": error,
    }


async def _generate_with_dalle(prompt_template: str, mood: str,
                                feedback: str, out_path: str,
                                *, api_key: str, model: str,
                                background_keyword: str = "") -> None:
    prompt = prompt_template
    if background_keyword:
        prompt = f"{prompt}, {background_keyword}"  # 사용자 직접 지정 keyword 우선 적용
    if mood:
        prompt = f"{prompt}, {mood} mood"
    if feedback:
        prompt = f"{prompt}. 추가 지시: {feedback}"
    prompt = f"{prompt}, no text, high quality"

    async with httpx.AsyncClient(timeout=DALLE_TIMEOUT_S) as client:
        resp = await client.post(
            "https://api.openai.com/v1/images/generations",
            headers={"Authorization": f"Bearer {api_key}"},
            json={"model": model, "prompt": prompt, "size": "1024x1024", "n": 1},
        )
        resp.raise_for_status()
        data = resp.json()["data"][0]
        if "url" in data:
            img_resp = await client.get(data["url"])
            img_resp.raise_for_status()
            img_bytes = img_resp.content
        elif "b64_json" in data:
            img_bytes = base64.b64decode(data["b64_json"])
        else:
            raise ValueError("DALL·E response has neither url nor b64_json")
        # PNG → JPG 변환
        with Image.open(BytesIO(img_bytes)) as src:
            img = src.convert("RGB")
        img.save(out_path, "JPEG", quality=92)