Update app.py

app.py

@@ -4,42 +4,57 @@ from PIL import Image
 import torch
 import requests
 import os
+from dotenv import load_dotenv
 
-# Load caption model
+# ───── 환경 변수 로딩 (토큰 안전하게 가져오기) ─────
+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"
+
+# ───── 이미지 캡셔닝 모델 로딩 ─────
 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")
 
-# Load ChatTTS (via inference API)
-CHAT_TTS_API = "https://api-inference.huggingface.co/models/2Noise/ChatTTS"
-headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
-
+# ───── 이미지 → 설명 문장 생성 ─────
 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)  # <- ✅ beam search 제거
+    output_ids = caption_model.generate(pixel_values, max_length=50)
     caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     return caption
 
-
-def tts_audio(text):
-    payload = {"inputs": text}
-    response = requests.post(CHAT_TTS_API, headers=headers, json=payload)
+# ───── 설명 → 음악 생성 (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"TTS API 오류: {response.status_code}, {response.text}")
-
+        raise Exception(f"MusicGen 오류: {response.status_code}, {response.text}")
 
+# ───── 전체 파이프라인 연결 ─────
 def process(image):
     caption = generate_caption(image)
-    audio = tts_audio(caption)
-    return caption, (audio, "result.wav")
+    prompt = f"A cheerful melody inspired by: {caption}"
+    audio = generate_music(prompt)
+    return caption, (audio, "musicgen_output.wav")
 
+# ───── Gradio 인터페이스 구성 ─────
 demo = gr.Interface(
     fn=process,
     inputs=gr.Image(type="pil"),
-    outputs=[gr.Text(label="설명"), gr.Audio(label="TTS 음성")],
-    title="🎨 AI 그림 설명 낭독기",
+    outputs=[
+        gr.Text(label="AI가 생성한 그림 설명"),
+        gr.Audio(label="생성된 AI 음악 (MusicGen)")
+    ],
+    title="🎨 AI 그림 음악 생성기",
+    description="그림을 업로드하면 AI가 설명을 만들고, 설명을 바탕으로 음악을 만들어 들려줍니다."
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()