drama-studio/audio-studio-musicgen/app/main.py

"""Audio Studio MusicGen API

AI 음악 생성 API 서버
"""

import logging
from contextlib import asynccontextmanager
from typing import Optional

from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.responses import Response
from pydantic import BaseModel, Field

from app.services.musicgen_service import musicgen_service

# 로깅 설정
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


# ========================================
# Pydantic 모델
# ========================================

class GenerateRequest(BaseModel):
    """음악 생성 요청"""
    prompt: str = Field(..., min_length=5, max_length=500, description="음악 설명")
    duration: int = Field(default=30, ge=5, le=30, description="생성 길이 (초)")
    top_k: int = Field(default=250, ge=50, le=500, description="top-k 샘플링")
    temperature: float = Field(default=1.0, ge=0.5, le=2.0, description="생성 다양성")


class HealthResponse(BaseModel):
    """헬스체크 응답"""
    status: str
    model_info: dict


# ========================================
# 앱 생명주기
# ========================================

@asynccontextmanager
async def lifespan(app: FastAPI):
    """앱 시작/종료 시 실행"""
    logger.info("MusicGen 서비스 시작...")
    try:
        await musicgen_service.initialize()
        logger.info("MusicGen 서비스 준비 완료")
    except Exception as e:
        logger.error(f"MusicGen 초기화 실패: {e}")
        # 초기화 실패해도 서버는 시작 (lazy loading 시도)

    yield

    logger.info("MusicGen 서비스 종료")


# ========================================
# FastAPI 앱
# ========================================

app = FastAPI(
    title="Audio Studio MusicGen",
    description="AI 음악 생성 API (Meta AudioCraft)",
    version="0.1.0",
    lifespan=lifespan,
)


# ========================================
# API 엔드포인트
# ========================================

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """헬스체크 엔드포인트"""
    return HealthResponse(
        status="healthy" if musicgen_service.is_initialized() else "initializing",
        model_info=musicgen_service.get_model_info(),
    )


@app.post("/generate")
async def generate_music(request: GenerateRequest):
    """텍스트 프롬프트로 음악 생성

    예시 프롬프트:
    - "upbeat electronic music for gaming"
    - "calm piano music, peaceful, ambient"
    - "energetic rock music with drums"
    - "lo-fi hip hop beats, relaxing"
    """
    try:
        audio_bytes = await musicgen_service.generate(
            prompt=request.prompt,
            duration=request.duration,
            top_k=request.top_k,
            temperature=request.temperature,
        )

        return Response(
            content=audio_bytes,
            media_type="audio/wav",
            headers={
                "X-Sample-Rate": "32000",
                "X-Duration": str(request.duration),
                "Content-Disposition": 'attachment; filename="generated_music.wav"',
            },
        )

    except Exception as e:
        logger.error(f"음악 생성 실패: {e}")
        raise HTTPException(status_code=500, detail=f"Music generation failed: {str(e)}")


@app.post("/generate-with-melody")
async def generate_with_melody(
    prompt: str = Form(..., min_length=5, description="음악 설명"),
    duration: int = Form(default=30, ge=5, le=30, description="생성 길이"),
    melody_audio: UploadFile = File(..., description="참조 멜로디 오디오"),
):
    """멜로디 조건부 음악 생성

    참조 멜로디의 멜로디/하모니를 유지하면서 새로운 음악 생성
    """
    try:
        melody_bytes = await melody_audio.read()

        if len(melody_bytes) < 1000:
            raise HTTPException(status_code=400, detail="Melody audio is too small")

        audio_bytes = await musicgen_service.generate_with_melody(
            prompt=prompt,
            melody_audio=melody_bytes,
            duration=duration,
        )

        return Response(
            content=audio_bytes,
            media_type="audio/wav",
            headers={
                "X-Sample-Rate": "32000",
                "X-Duration": str(duration),
                "Content-Disposition": 'attachment; filename="melody_based_music.wav"',
            },
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"멜로디 기반 생성 실패: {e}")
        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")


@app.get("/prompts")
async def get_example_prompts():
    """예시 프롬프트 목록"""
    return {
        "examples": [
            {
                "category": "Electronic",
                "prompts": [
                    "upbeat electronic dance music with synthesizers",
                    "chill electronic ambient music",
                    "retro synthwave 80s style music",
                ],
            },
            {
                "category": "Classical",
                "prompts": [
                    "calm piano solo, classical style",
                    "orchestral epic cinematic music",
                    "gentle string quartet, romantic",
                ],
            },
            {
                "category": "Pop/Rock",
                "prompts": [
                    "energetic rock music with electric guitar",
                    "upbeat pop song with catchy melody",
                    "acoustic guitar folk music",
                ],
            },
            {
                "category": "Ambient/Lo-fi",
                "prompts": [
                    "lo-fi hip hop beats, relaxing, study music",
                    "peaceful ambient nature sounds music",
                    "meditation music, calm, zen",
                ],
            },
            {
                "category": "Game/Film",
                "prompts": [
                    "epic adventure game soundtrack",
                    "tense suspenseful thriller music",
                    "cheerful happy video game background",
                ],
            },
        ]
    }