Update app.py

app.py

@@ -1,88 +1,78 @@
-import os
-import sys
-import types
-import subprocess
-import tempfile
-import torch
-import gradio as gr
-from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+import os, sys, types, subprocess, tempfile
+import torch, gradio as gr
+from transformers import (
+    VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+)
 from PIL import Image
 
 # ── 환경 변수 ────────────────────────────────────────────────
 os.environ["HF_FORCE_SAFE_SERIALIZATION"] = "1"
-os.environ["XFORMERS_FORCE_DISABLE"] = "1"      # 실제 xformers 비활성화
+os.environ["XFORMERS_FORCE_DISABLE"] = "1"    # audiocraft 내부 플래그
 
-# ── ✨ xformers 더미 모듈 삽입 ─────────────────────────────────
+# ── xformers 더미 모듈 ───────────────────────────────────────
 dummy = types.ModuleType("xformers")
 dummy.__version__ = "0.0.0"
-
-# 하위 모듈 xformers.ops
 ops = types.ModuleType("xformers.ops")
 
-def _fake_memory_efficient_attention(q, k, v, *_, dropout_p: float = 0.0, **__):
-    """
-    xformers.memory_efficient_attention 대체 구현.
-    PyTorch 2.x 기본 S-DPA로 처리해 속도는 느리지만 CPU에서도 동작합니다.
-    """
+def _fake_mem_eff_attn(q, k, v, *_, dropout_p: float = 0.0, **__):
     return torch.nn.functional.scaled_dot_product_attention(
         q, k, v, dropout_p=dropout_p, is_causal=False
     )
+class _FakeLowerTriangularMask: pass
 
-class _FakeLowerTriangularMask:  # audiocraft 내부 타입 체크용 더미
-    pass
-
-ops.memory_efficient_attention = _fake_memory_efficient_attention
+ops.memory_efficient_attention = _fake_mem_eff_attn
 ops.LowerTriangularMask = _FakeLowerTriangularMask
-
 dummy.ops = ops
 sys.modules["xformers"] = dummy
 sys.modules["xformers.ops"] = ops
 # ────────────────────────────────────────────────────────────
 
-# ── audiocraft 동적 설치 ─────────────────────────────────────
+# ── audiocraft 로드 (postInstall에서 이미 설치됐음) ───────────
 try:
     from audiocraft.models import MusicGen
-except ModuleNotFoundError:
+    from audiocraft.data.audio import audio_write
+except ModuleNotFoundError:                 # 예외적 로컬 실행 대비
     subprocess.check_call([
         sys.executable, "-m", "pip", "install",
         "git+https://github.com/facebookresearch/audiocraft@main",
-        "--use-pep517"
+        "--no-deps", "--use-pep517"
     ])
     from audiocraft.models import MusicGen
-
-from audiocraft.data.audio import audio_write
+    from audiocraft.data.audio import audio_write
 
 # ── 이미지 캡셔닝 모델 ─────────────────────────────────────
 caption_model = VisionEncoderDecoderModel.from_pretrained(
     "nlpconnect/vit-gpt2-image-captioning",
-    use_safetensors=True,
-    low_cpu_mem_usage=True
+    use_safetensors=True, low_cpu_mem_usage=True
+)
+feature_extractor = ViTImageProcessor.from_pretrained(
+    "nlpconnect/vit-gpt2-image-captioning"
+)
+tokenizer = AutoTokenizer.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 = MusicGen.get_pretrained("facebook/musicgen-small")
-musicgen.set_generation_params(duration=10)  # 10초 음악
+musicgen.set_generation_params(duration=10)
 
-# ── 유틸 함수들 ─────────────────────────────────────────────
+# ── 파이프라인 함수들 ──────────────────────────────────────
 def generate_caption(image: Image.Image) -> str:
     pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
-    output_ids = caption_model.generate(pixel_values, max_length=50)
-    return tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    ids = caption_model.generate(pixel_values, max_length=50)
+    return tokenizer.decode(ids[0], skip_special_tokens=True)
 
 def generate_music(prompt: str) -> str:
-    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
+    wav = musicgen.generate([prompt])
+    tmpdir = tempfile.mkdtemp()
+    path = os.path.join(tmpdir, "musicgen.wav")
+    audio_write(path, wav[0], musicgen.sample_rate, strategy="loudness")
+    return path
 
 def process(image: Image.Image):
     caption = generate_caption(image)
-    prompt = f"A cheerful melody inspired by: {caption}"
-    audio_path = generate_music(prompt)
-    return caption, audio_path
+    path = generate_music(f"A cheerful melody inspired by: {caption}")
+    return caption, path
 
 # ── Gradio UI ──────────────────────────────────────────────
 demo = gr.Interface(
@@ -93,7 +83,7 @@ demo = gr.Interface(
         gr.Audio(label="생성된 AI 음악 (MusicGen)")
     ],
     title="🎨 AI 그림‑음악 생성기",
-    description="그림을 업로드하면 AI가 설명을 만들고, 설명을 바탕으로 음악을 생성해 들려줍니다."
+    description="그림을 업로드하면 AI가 설명을 만들고, 설명을 바탕으로 음악을 10초간 생성해 들려줍니다."
 )
 
 if __name__ == "__main__":