Update app.py

app.py

@@ -1,114 +1,57 @@
-import os, sys, types, subprocess, tempfile
-import torch, gradio as gr
-from transformers import (
-    VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
-)
+import os, io, base64, tempfile, requests
+import gradio as gr
 from PIL import Image
 
 # ─────────────────────────────────────────────────────────────
-# 0. 환경 변수
-# ─────────────────────────────────────────────────────────────
-os.environ["HF_FORCE_SAFE_SERIALIZATION"] = "1"
-os.environ["XFORMERS_FORCE_DISABLE"] = "1"    # audiocraft 내부 플래그
-
+# 1. 환경 변수 & HF Inference API 설정
 # ─────────────────────────────────────────────────────────────
-# 1. xformers 더미 모듈 주입 (GPU 종속 제거)
-# ─────────────────────────────────────────────────────────────
-dummy = types.ModuleType("xformers")
-dummy.__version__ = "0.0.0"
-ops = types.ModuleType("xformers.ops")
-
-def _fake_mea(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
-    )
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise RuntimeError("HF_TOKEN 비밀 값이 설정되어 있지 않습니다. Spaces Settings → Secrets에서 등록해 주세요.")
 
-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
+HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
+CAPTION_API = "https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-base"
+MUSIC_API   = "https://api-inference.huggingface.co/models/facebook/musicgen-small"
 
 # ─────────────────────────────────────────────────────────────
-# 2. (선택) 설치하지 않은 모듈만 안전망으로 스텁 처리 ★
-#    - 이미 requirements.txt에서 설치한 모듈(librosa, av 등)은
-#      스텁 대상에서 제거합니다.
+# 2. 이미지 캡션 생성 (BLIP-base via API)
 # ─────────────────────────────────────────────────────────────
-for name in ("pesq", "pystoi", "soxr"):   # ★ 필요시만 남김
-    if name not in sys.modules:
-        sys.modules[name] = types.ModuleType(name)
+def generate_caption(image_pil: Image.Image) -> str:
+    buf = io.BytesIO()
+    image_pil.save(buf, format="PNG")
+    buf.seek(0)
 
-# ─────────────────────────────────────────────────────────────
-# 3. audiocraft (MusicGen) 불러오기
-# ─────────────────────────────────────────────────────────────
-try:
-    from audiocraft.models import MusicGen
-    from audiocraft.data.audio import audio_write
-except ModuleNotFoundError:
-    subprocess.check_call([
-        sys.executable, "-m", "pip", "install",
-        "git+https://github.com/facebookresearch/audiocraft@main",
-        "--no-deps", "--use-pep517"
-    ])
-    subprocess.check_call([sys.executable, "-m", "pip", "install",
-                           "encodec", "librosa", "av", "torchdiffeq",
-                           "torchmetrics", "num2words"])
-    from audiocraft.models import MusicGen
-    from audiocraft.data.audio import audio_write
+    # binary upload 방식
+    response = requests.post(CAPTION_API, headers=HEADERS, data=buf.getvalue(), timeout=60)
+    response.raise_for_status()
+    result = response.json()
+    # API 응답: [{"generated_text": "..."}]
+    return result[0]["generated_text"]
 
 # ─────────────────────────────────────────────────────────────
-# 4. 이미지 캡셔닝 모델
+# 3. MusicGen-small 음악 생성 (10초, via API)
 # ─────────────────────────────────────────────────────────────
-# 4. 이미지 캡셔닝 모델 ------------------------------------
-caption_model = VisionEncoderDecoderModel.from_pretrained(
-    "nlpconnect/vit-gpt2-image-captioning",
-    use_safetensors=True,        # 그대로
-    low_cpu_mem_usage=False,     # ← meta 로딩 비활성화
-    device_map=None              # ← Accelerate 자동 분할 끄기
-).eval()                         # 평가 모드
-
-feature_extractor = ViTImageProcessor.from_pretrained(
-    "nlpconnect/vit-gpt2-image-captioning"
-)
-tokenizer = AutoTokenizer.from_pretrained(
-    "nlpconnect/vit-gpt2-image-captioning"
-)
-
+def generate_music(prompt: str, duration: int = 10) -> str:
+    payload = {"inputs": prompt, "parameters": {"duration": duration}}
+    response = requests.post(MUSIC_API, headers=HEADERS, json=payload, timeout=120)
+    response.raise_for_status()
 
-# ─────────────────────────────────────────────────────────────
-# 5. MusicGen 모델
-# ─────────────────────────────────────────────────────────────
-musicgen = MusicGen.get_pretrained("facebook/musicgen-small")
-musicgen.set_generation_params(duration=10)
+    # API 응답은 WAV 바이너리
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+    tmp.write(response.content)
+    tmp.close()
+    return tmp.name
 
 # ─────────────────────────────────────────────────────────────
-# 6. 파이프라인 함수
+# 4. 전체 파이프라인
 # ─────────────────────────────────────────────────────────────
-def generate_caption(image: Image.Image) -> str:
-    with torch.no_grad():
-        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)
-
-
-
-def generate_music(prompt: str) -> str:
-    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")
-    return path
-
-def process(image: Image.Image):
+def process(image):
     caption = generate_caption(image)
-    path    = generate_music(f"A cheerful melody inspired by: {caption}")
-    return caption, path
+    audio   = generate_music(f"A cheerful melody inspired by: {caption}")
+    return caption, audio
 
 # ─────────────────────────────────────────────────────────────
-# 7. Gradio UI
+# 5. Gradio 인터페이스
 # ─────────────────────────────────────────────────────────────
 demo = gr.Interface(
     fn=process,
@@ -117,8 +60,10 @@ demo = gr.Interface(
         gr.Text(label="AI가 생성한 그림 설명"),
         gr.Audio(label="생성된 AI 음악 (MusicGen)")
     ],
-    title="🎨 AI 그림-음악 생성기",
-    description="그림을 업로드하면 AI가 설명을 만들고, 설명을 바탕으로 10초 길이의 음악을 생성해 들려줍니다."
+    title="🎨 AI 그림-음악 생성기 (Inference API 버전)",
+    description="이미지를 업로드하면 BLIP-base가 설명을 생성하고, 해당 설명으로 MusicGen-small이 10초 음악을 만듭니다.",
+    concurrency_count=1,          # 메모리 보호용: 동시 1요청
+    cache_examples=False
 )
 
 if __name__ == "__main__":