Update app.py

app.py

@@ -1,26 +1,21 @@
 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 requests
 import os
-from dotenv import load_dotenv
-
-# ───── 환경 변수 로딩 (토큰 안전하게 가져오기) ─────
-load_dotenv()
-HF_TOKEN = os.getenv("HF_TOKEN")
-if HF_TOKEN is None:
-    raise ValueError("HF_TOKEN이 .env 파일에 없습니다.")
-
-# ───── Hugging Face Inference API 설정 ─────
-headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-MUSICGEN_API = "https://api-inference.huggingface.co/models/facebook/musicgen-small"
+import tempfile
 
 # ───── 이미지 캡셔닝 모델 로딩 ─────
 caption_model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 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
@@ -28,21 +23,20 @@ def generate_caption(image):
     caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     return caption
 
-# ───── 설명 → 음악 생성 (MusicGen API 호출) ─────
+# ───── 설명 → 음악 생성 ─────
 def generate_music(prompt):
-    payload = {"inputs": prompt}
-    response = requests.post(MUSICGEN_API, headers=headers, json=payload)
-    if response.status_code == 200:
-        return response.content
-    else:
-        raise Exception(f"MusicGen 오류: {response.status_code}, {response.text}")
+    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 = generate_music(prompt)
-    return caption, (audio, "musicgen_output.wav")
+    audio_path = generate_music(prompt)
+    return caption, audio_path
 
 # ───── Gradio 인터페이스 구성 ─────
 demo = gr.Interface(