Create app.py

app.py

@@ -0,0 +1,354 @@
+import spaces
+import logging
+from datetime import datetime
+from pathlib import Path
+import gradio as gr
+import torch
+import torchaudio
+import os
+import requests
+from transformers import pipeline
+import tempfile
+import numpy as np
+from einops import rearrange
+import cv2
+from scipy.io import wavfile
+import librosa
+import json
+from typing import Optional, Tuple, List
+import atexit
+
+# 환경 변수 설정으로 torch.load 체크 우회 (임시 해결책)
+os.environ["TRANSFORMERS_ALLOW_UNSAFE_DESERIALIZATION"] = "1"
+
+try:
+    import mmaudio
+except ImportError:
+    os.system("pip install -e .")
+    import mmaudio
+
+from mmaudio.eval_utils import (ModelConfig, all_model_cfg, generate, load_video, make_video,
+                                setup_eval_logging)
+from mmaudio.model.flow_matching import FlowMatching
+from mmaudio.model.networks import MMAudio, get_my_mmaudio
+from mmaudio.model.sequence_config import SequenceConfig
+from mmaudio.model.utils.features_utils import FeaturesUtils
+
+# 로깅 설정
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+log = logging.getLogger()
+
+# CUDA 설정
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+else:
+    device = torch.device("cpu")
+
+dtype = torch.bfloat16
+
+# 모델 설정
+model: ModelConfig = all_model_cfg['large_44k_v2']
+model.download_if_needed()
+output_dir = Path('./output/gradio')
+
+setup_eval_logging()
+
+# 번역기 설정 - safetensors 사용 시도
+try:
+    # 먼저 safetensors 형식이 있는지 확인
+    translator = pipeline("translation", 
+                         model="Helsinki-NLP/opus-mt-ko-en", 
+                         device="cpu",
+                         use_fast=True,  # Fast tokenizer 사용
+                         trust_remote_code=False)
+except Exception as e:
+    log.warning(f"Failed to load translation model with safetensors: {e}")
+    # 대체 방법: 환경 변수 설정 후 로드
+    try:
+        translator = pipeline("translation", 
+                             model="Helsinki-NLP/opus-mt-ko-en", 
+                             device="cpu")
+    except Exception as e2:
+        log.error(f"Failed to load translation model: {e2}")
+        translator = None
+
+PIXABAY_API_KEY = "33492762-a28a596ec4f286f84cd328b17"
+
+def cleanup_temp_files():
+    temp_dir = tempfile.gettempdir()
+    for file in os.listdir(temp_dir):
+        if file.endswith(('.mp4', '.flac')):
+            try:
+                os.remove(os.path.join(temp_dir, file))
+            except:
+                pass
+
+atexit.register(cleanup_temp_files)
+
+def get_model() -> tuple[MMAudio, FeaturesUtils, SequenceConfig]:
+    with torch.cuda.device(device):
+        seq_cfg = model.seq_cfg
+        net: MMAudio = get_my_mmaudio(model.model_name).to(device, dtype).eval()
+        net.load_weights(torch.load(model.model_path, map_location=device, weights_only=True))
+        log.info(f'Loaded weights from {model.model_path}')
+
+        feature_utils = FeaturesUtils(
+            tod_vae_ckpt=model.vae_path,
+            synchformer_ckpt=model.synchformer_ckpt,
+            enable_conditions=True,
+            mode=model.mode,
+            bigvgan_vocoder_ckpt=model.bigvgan_16k_path,
+            need_vae_encoder=False
+        ).to(device, dtype).eval()
+
+        return net, feature_utils, seq_cfg
+
+net, feature_utils, seq_cfg = get_model()
+
+# translate_prompt 함수 수정
+def translate_prompt(text):
+    try:
+        # 번역기가 없으면 원본 텍스트 반환
+        if translator is None:
+            return text
+            
+        if text and any(ord(char) >= 0x3131 and ord(char) <= 0xD7A3 for char in text):
+            # CPU에서 번역 실행
+            with torch.no_grad():
+                translation = translator(text)[0]['translation_text']
+            return translation
+        return text
+    except Exception as e:
+        logging.error(f"Translation error: {e}")
+        return text
+
+# search_videos 함수 수정
+@torch.no_grad()
+def search_videos(query):
+    try:
+        # CPU에서 번역 실행
+        query = translate_prompt(query)
+        return search_pixabay_videos(query, PIXABAY_API_KEY)
+    except Exception as e:
+        logging.error(f"Video search error: {e}")
+        return []
+
+def search_pixabay_videos(query, api_key):
+    try:
+        base_url = "https://pixabay.com/api/videos/"
+        params = {
+            "key": api_key,
+            "q": query,
+            "per_page": 40
+        }
+        
+        response = requests.get(base_url, params=params)
+        if response.status_code == 200:
+            data = response.json()
+            return [video['videos']['large']['url'] for video in data.get('hits', [])]
+        return []
+    except Exception as e:
+        logging.error(f"Pixabay API error: {e}")
+        return []
+
+@spaces.GPU
+@torch.inference_mode()
+def video_to_audio(video: gr.Video, prompt: str, negative_prompt: str, seed: int, num_steps: int,
+                   cfg_strength: float, duration: float):
+    prompt = translate_prompt(prompt)
+    negative_prompt = translate_prompt(negative_prompt)
+
+    rng = torch.Generator(device=device)
+    rng.manual_seed(seed)
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
+
+    clip_frames, sync_frames, duration = load_video(video, duration)
+    clip_frames = clip_frames.unsqueeze(0)
+    sync_frames = sync_frames.unsqueeze(0)
+    seq_cfg.duration = duration
+    net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+    audios = generate(clip_frames,
+                      sync_frames, [prompt],
+                      negative_text=[negative_prompt],
+                      feature_utils=feature_utils,
+                      net=net,
+                      fm=fm,
+                      rng=rng,
+                      cfg_strength=cfg_strength)
+    audio = audios.float().cpu()[0]
+
+    video_save_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
+    make_video(video,
+               video_save_path,
+               audio,
+               sampling_rate=seq_cfg.sampling_rate,
+               duration_sec=seq_cfg.duration)
+    return video_save_path
+
+@spaces.GPU
+@torch.inference_mode()
+def text_to_audio(prompt: str, negative_prompt: str, seed: int, num_steps: int, cfg_strength: float,
+                  duration: float):
+    prompt = translate_prompt(prompt)
+    negative_prompt = translate_prompt(negative_prompt)
+
+    rng = torch.Generator(device=device)
+    rng.manual_seed(seed)
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
+
+    clip_frames = sync_frames = None
+    seq_cfg.duration = duration
+    net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+    audios = generate(clip_frames,
+                      sync_frames, [prompt],
+                      negative_text=[negative_prompt],
+                      feature_utils=feature_utils,
+                      net=net,
+                      fm=fm,
+                      rng=rng,
+                      cfg_strength=cfg_strength)
+    audio = audios.float().cpu()[0]
+
+    audio_save_path = tempfile.NamedTemporaryFile(delete=False, suffix='.flac').name
+    torchaudio.save(audio_save_path, audio, seq_cfg.sampling_rate)
+    return audio_save_path
+
+# CSS 스타일
+custom_css = """
+.gradio-container {
+    background: linear-gradient(45deg, #1a1a1a, #2a2a2a);
+    border-radius: 15px;
+    box-shadow: 0 8px 32px rgba(0,0,0,0.3);
+    color: #e0e0e0;
+}
+
+.input-container, .output-container {
+    background: rgba(40, 40, 40, 0.95);
+    backdrop-filter: blur(10px);
+    border-radius: 10px;
+    padding: 20px;
+    transform-style: preserve-3d;
+    transition: transform 0.3s ease;
+    border: 1px solid rgba(255, 255, 255, 0.1);
+}
+
+.input-container:hover {
+    transform: translateZ(20px);
+    box-shadow: 0 8px 32px rgba(0,0,0,0.5);
+}
+
+.gallery-item {
+    transition: transform 0.3s ease;
+    border-radius: 8px;
+    overflow: hidden;
+    background: #2a2a2a;
+}
+
+.gallery-item:hover {
+    transform: scale(1.05);
+    box-shadow: 0 4px 15px rgba(0,0,0,0.4);
+}
+
+.tabs {
+    background: rgba(30, 30, 30, 0.95);
+    border-radius: 10px;
+    padding: 10px;
+    border: 1px solid rgba(255, 255, 255, 0.05);
+}
+
+button {
+    background: linear-gradient(45deg, #2196F3, #1976D2);
+    border: none;
+    border-radius: 5px;
+    transition: all 0.3s ease;
+    color: white;
+}
+
+button:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 4px 15px rgba(33,150,243,0.3);
+}
+
+textarea, input[type="text"], input[type="number"] {
+    background: rgba(30, 30, 30, 0.95) !important;
+    border: 1px solid rgba(255, 255, 255, 0.1) !important;
+    color: #e0e0e0 !important;
+    border-radius: 5px !important;
+}
+
+label {
+    color: #e0e0e0 !important;
+}
+
+.gallery {
+    background: rgba(30, 30, 30, 0.95);
+    padding: 15px;
+    border-radius: 10px;
+    border: 1px solid rgba(255, 255, 255, 0.05);
+}
+"""
+
+css = """
+footer {
+    visibility: hidden;
+}
+""" + custom_css
+
+# Gradio 인터페이스 생성
+text_to_audio_tab = gr.Interface(
+    fn=text_to_audio,
+    inputs=[
+        gr.Textbox(label="Prompt(한글지원)" if translator else "Prompt"),
+        gr.Textbox(label="Negative Prompt"),
+        gr.Number(label="Seed", value=0),
+        gr.Number(label="Steps", value=25),
+        gr.Number(label="Guidance Scale", value=4.5),
+        gr.Number(label="Duration (sec)", value=8),
+    ],
+    outputs=gr.Audio(label="Generated Audio"),
+    css=custom_css
+)
+
+video_to_audio_tab = gr.Interface(
+    fn=video_to_audio,
+    inputs=[
+        gr.Video(label="Input Video"),
+        gr.Textbox(label="Prompt(한글지원)" if translator else "Prompt"),
+        gr.Textbox(label="Negative Prompt", value="music"),
+        gr.Number(label="Seed", value=0),
+        gr.Number(label="Steps", value=25),
+        gr.Number(label="Guidance Scale", value=4.5),
+        gr.Number(label="Duration (sec)", value=8),
+    ],
+    outputs=gr.Video(label="Generated Result"),
+    css=custom_css
+)
+
+video_search_tab = gr.Interface(
+    fn=search_videos,
+    inputs=gr.Textbox(label="Search Query(한글지원)" if translator else "Search Query"),
+    outputs=gr.Gallery(label="Search Results", columns=4, rows=20),
+    css=custom_css,
+    api_name=False
+)
+
+# 메인 실행
+if __name__ == "__main__":
+    # 번역기 로드 실패 시 경고 메시지
+    if translator is None:
+        log.warning("Translation model failed to load. Korean translation will be disabled.")
+        
+    gr.TabbedInterface(
+        [video_search_tab, video_to_audio_tab, text_to_audio_tab],
+        ["Video Search", "Video-to-Audio", "Text-to-Audio"],
+        theme="soft",
+        css=css
+    ).launch(allowed_paths=[output_dir])