import os, tempfile, soundfile as sf
import gradio as gr
from PIL import Image
from transformers import pipeline

# ────────────────────────────────────────────────
# 1. 파이프라인 로드 (CPU: device=-1)
# ────────────────────────────────────────────────
CAPTION_ID = "Salesforce/blip-image-captioning-base"   # 용량↓: blip-image-captioning
MUSIC_ID   = "facebook/musicgen-melody"                # 용량↓: musicgen-small

caption_pipe = pipeline(
    "image-to-text",
    model=CAPTION_ID,
    device=-1
)

music_pipe = pipeline(
    "text-to-audio",
    model=MUSIC_ID,
    device=-1,
    generate_kwargs={"duration": 10}   # 초 단위
)

# ────────────────────────────────────────────────
# 2. 유틸 함수
# ────────────────────────────────────────────────
def generate_caption(img: Image.Image) -> str:
    return caption_pipe(img)[0]["generated_text"]

def generate_music(prompt: str) -> str:
    result = music_pipe(prompt, forward_params={"do_sample": True})[0]
    audio, sr = result["audio"], result["sampling_rate"]

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    sf.write(tmp.name, audio, sr)
    return tmp.name

# ────────────────────────────────────────────────
# 3. 전체 파이프라인
# ────────────────────────────────────────────────
def process(image):
    caption = generate_caption(image)
    audio   = generate_music(f"A cheerful melody inspired by: {caption}")
    return caption, audio

# ────────────────────────────────────────────────
# 4. Gradio UI
# ────────────────────────────────────────────────
demo = gr.Interface(
    fn=process,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Text(label="AI가 생성한 그림 설명"),
        gr.Audio(label="생성된 AI 음악 (MusicGen)")
    ],
    title="🎨 로컬 BLIP-base + MusicGen-melody",
    description="이미지를 업로드하면 BLIP-base가 설명을 생성하고, "
                "그 설명으로 MusicGen-melody가 10초 음악을 만듭니다."
).queue()

if __name__ == "__main__":
    demo.launch()