import gradio as gr
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
from PIL import Image
import torch
import os
import tempfile

# ───── 이미지 캡셔닝 모델 로딩 ─────
caption_model = VisionEncoderDecoderModel.from_pretrained(
    "nlpconnect/vit-gpt2-image-captioning", 
    safetensors=True  # ✅ torch.load() 대신 안전한 방식 사용
)
feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")

# ───── MusicGen 모델 로딩 ─────
musicgen = MusicGen.get_pretrained("facebook/musicgen-small")
musicgen.set_generation_params(duration=10)  # 생성할 음악 길이 (초)

# ───── 이미지 → 설명 문장 생성 ─────
def generate_caption(image):
    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
    output_ids = caption_model.generate(pixel_values, max_length=50)
    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return caption

# ───── 설명 → 음악 생성 ─────
def generate_music(prompt):
    wav = musicgen.generate([prompt])  # batch size 1
    tmp_dir = tempfile.mkdtemp()
    audio_path = os.path.join(tmp_dir, "musicgen_output.wav")
    audio_write(audio_path, wav[0], musicgen.sample_rate, strategy="loudness")
    return audio_path

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

# ───── Gradio 인터페이스 구성 ─────
demo = gr.Interface(
    fn=process,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Text(label="AI가 생성한 그림 설명"),
        gr.Audio(label="생성된 AI 음악 (MusicGen)")
    ],
    title="🎨 AI 그림 음악 생성기",
    description="그림을 업로드하면 AI가 설명을 만들고, 설명을 바탕으로 음악을 만들어 들려줍니다."
)

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