Update app.py

app.py

@@ -5,45 +5,73 @@ from transformers import (
 )
 from PIL import Image
 
-# ── 환경 변수 ────────────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
+# 0. 환경 변수
+# ─────────────────────────────────────────────────────────────
 os.environ["HF_FORCE_SAFE_SERIALIZATION"] = "1"
 os.environ["XFORMERS_FORCE_DISABLE"] = "1"    # audiocraft 내부 플래그
 
-# ── xformers 더미 모듈 ───────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
+# 1. xformers 더미 모듈 주입 (GPU 종속 제거)
+# ─────────────────────────────────────────────────────────────
 dummy = types.ModuleType("xformers")
 dummy.__version__ = "0.0.0"
 ops = types.ModuleType("xformers.ops")
 
-def _fake_mem_eff_attn(q, k, v, *_, dropout_p: float = 0.0, **__):
+def _fake_mea(q, k, v, *_, dropout_p: float = 0.0, **__):
+    # PyTorch 2.x 표준 S-DPA로 대체 (CPU에서도 동작)
     return torch.nn.functional.scaled_dot_product_attention(
         q, k, v, dropout_p=dropout_p, is_causal=False
     )
-class _FakeLowerTriangularMask: pass
 
-ops.memory_efficient_attention = _fake_mem_eff_attn
+class _FakeLowerTriangularMask:
+    """audiocraft가 타입 존재만 확인하므로 빈 클래스로 대체"""
+    pass
+
+ops.memory_efficient_attention = _fake_mea
 ops.LowerTriangularMask = _FakeLowerTriangularMask
 dummy.ops = ops
 sys.modules["xformers"] = dummy
 sys.modules["xformers.ops"] = ops
-# ────────────────────────────────────────────────────────────
 
-# ── audiocraft 로드 (postInstall에서 이미 설치됐음) ───────────
+# ─────────────────────────────────────────────────────────────
+# 2. 기타 누락 모듈에 대한 더미(stub) 안전망
+#    (requirements.txt 에 이미 설치하지만, 혹시 빠져도 런타임 통과)
+# ─────────────────────────────────────────────────────────────
+for name in (
+    "av", "librosa", "torchdiffeq", "torchmetrics",
+    "pesq", "pystoi", "soxr"
+):
+    if name not in sys.modules:
+        sys.modules[name] = types.ModuleType(name)
+
+# ─────────────────────────────────────────────────────────────
+# 3. audiocraft (MusicGen) 불러오기
+# ─────────────────────────────────────────────────────────────
 try:
     from audiocraft.models import MusicGen
     from audiocraft.data.audio import audio_write
-except ModuleNotFoundError:                 # 예외적 로컬 실행 대비
+except ModuleNotFoundError:
+    # 로컬 실행 등으로 미설치 시: 의존성 없는 형태로 설치
     subprocess.check_call([
         sys.executable, "-m", "pip", "install",
         "git+https://github.com/facebookresearch/audiocraft@main",
         "--no-deps", "--use-pep517"
     ])
+    # 필요 최소 의존성만 즉석 설치 (stub로도 대부분 통과하지만 안전하게)
+    subprocess.check_call([sys.executable, "-m", "pip", "install",
+                           "encodec", "torchdiffeq", "torchmetrics",
+                           "librosa", "soxr", "av"])
     from audiocraft.models import MusicGen
     from audiocraft.data.audio import audio_write
 
-# ── 이미지 캡셔닝 모델 ─────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
+# 4. 이미지 캡셔닝 모델
+# ─────────────────────────────────────────────────────────────
 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"
@@ -52,18 +80,22 @@ tokenizer = AutoTokenizer.from_pretrained(
     "nlpconnect/vit-gpt2-image-captioning"
 )
 
-# ── MusicGen 모델 ──────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
+# 5. MusicGen 모델 (CPU 전용)
+# ─────────────────────────────────────────────────────────────
 musicgen = MusicGen.get_pretrained("facebook/musicgen-small")
-musicgen.set_generation_params(duration=10)
+musicgen.set_generation_params(duration=10)  # 10초 길이
 
-# ── 파이프라인 함수들 ──────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
+# 6. 파이프라인 함수
+# ─────────────────────────────────────────────────────────────
 def generate_caption(image: Image.Image) -> str:
     pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
     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])
+    wav = musicgen.generate([prompt])           # batch size = 1
     tmpdir = tempfile.mkdtemp()
     path = os.path.join(tmpdir, "musicgen.wav")
     audio_write(path, wav[0], musicgen.sample_rate, strategy="loudness")
@@ -71,10 +103,12 @@ def generate_music(prompt: str) -> str:
 
 def process(image: Image.Image):
     caption = generate_caption(image)
-    path = generate_music(f"A cheerful melody inspired by: {caption}")
+    path    = generate_music(f"A cheerful melody inspired by: {caption}")
     return caption, path
 
-# ── Gradio UI ──────────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────
+# 7. Gradio UI
+# ─────────────────────────────────────────────────────────────
 demo = gr.Interface(
     fn=process,
     inputs=gr.Image(type="pil"),
@@ -82,8 +116,8 @@ demo = gr.Interface(
         gr.Text(label="AI가 생성한 그림 설명"),
         gr.Audio(label="생성된 AI 음악 (MusicGen)")
     ],
-    title="🎨 AI 그림‑음악 생성기",
-    description="그림을 업로드하면 AI가 설명을 만들고, 설명을 바탕으로 음악을 10초간 생성해 들려줍니다."
+    title="🎨 AI 그림-음악 생성기",
+    description="그림을 업로드하면 AI가 설명을 만들고, 설명을 바탕으로 10초 길이의 음악을 생성해 들려줍니다."
 )
 
 if __name__ == "__main__":