app.py

import os
import sys
import types
import subprocess
import tempfile
import torch
import 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 비활성화

# ── ✨ 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에서도 동작합니다.
    """
    return torch.nn.functional.scaled_dot_product_attention(
        q, k, v, dropout_p=dropout_p, is_causal=False
    )

class _FakeLowerTriangularMask:  # audiocraft 내부 타입 체크용 더미
    pass

ops.memory_efficient_attention = _fake_memory_efficient_attention
ops.LowerTriangularMask = _FakeLowerTriangularMask

dummy.ops = ops
sys.modules["xformers"] = dummy
sys.modules["xformers.ops"] = ops
# ────────────────────────────────────────────────────────────

# ── audiocraft 동적 설치 ─────────────────────────────────────
try:
    from audiocraft.models import MusicGen
except ModuleNotFoundError:
    subprocess.check_call([
        sys.executable, "-m", "pip", "install",
        "git+https://github.com/facebookresearch/audiocraft@main",
        "--use-pep517"
    ])
    from audiocraft.models import MusicGen

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
)
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)  # 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)

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

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

# ── Gradio UI ──────────────────────────────────────────────
demo = gr.Interface(
    fn=process,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Text(label="AI가 생성한 그림 설명"),
        gr.Audio(label="생성된 AI 음악 (MusicGen)")
    ],
    title="🎨 AI 그림‑음악 생성기",
    description="그림을 업로드하면 AI가 설명을 만들고, 설명을 바탕으로 음악을 생성해 들려줍니다."
)

if __name__ == "__main__":
    demo.launch()