Update app.py

app.py

@@ -1,3 +1,160 @@
+import spaces
+import gc
+import logging
+import sys
+import os
+from datetime import datetime
+from fractions import Fraction
+from pathlib import Path
+
+import gradio as gr
+import torch
+import torchaudio
+from huggingface_hub import hf_hub_download
+
+# --- CONFIGURATION & PATHS ---
+MY_VAULT_REPO = "ibyteohdear/mmaudio-weights-vault" 
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(current_dir, "MMAudio"))
+
+from mmaudio.eval_utils import (ModelConfig, VideoInfo, all_model_cfg, generate, load_image,
+                                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
+
+# Optimization flags
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+log = logging.getLogger()
+setup_eval_logging()
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
+output_dir = Path('./output/gradio')
+
+# --- WEIGHT SYNCHRONIZATION ---
+def get_weights():
+    log.info(f"Syncing weights from {MY_VAULT_REPO}...")
+    
+    return {
+        "model": hf_hub_download(repo_id=MY_VAULT_REPO, filename="weights/mmaudio_large_44k_v2.pth", token=HF_TOKEN),
+        "vae": hf_hub_download(repo_id=MY_VAULT_REPO, filename="ext_weights/v1-44.pth", token=HF_TOKEN), 
+        "sync": hf_hub_download(repo_id=MY_VAULT_REPO, filename="ext_weights/synchformer_state_dict.pth", token=HF_TOKEN),
+        "vocoder": hf_hub_download(repo_id=MY_VAULT_REPO, filename="ext_weights/v1-44.pth", token=HF_TOKEN)
+    }
+
+weight_paths = get_weights()
+
+# --- MODEL INITIALIZATION ---
+model: ModelConfig = all_model_cfg['large_44k_v2']
+model.model_path = weight_paths["model"]
+model.vae_path = weight_paths["vae"]
+model.synchformer_ckpt = weight_paths["sync"]
+model.bigvgan_16k_path = weight_paths["vocoder"]
+
+def load_all_models() -> tuple[MMAudio, FeaturesUtils, SequenceConfig]:
+    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))
+    
+    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 = load_all_models()
+
+# --- INFERENCE FUNCTIONS ---
+
+@spaces.GPU()
+@torch.inference_mode()
+def video_to_audio(video, prompt, negative_prompt, seed, num_steps, cfg_strength, duration):
+    rng = torch.Generator(device=device)
+    rng.manual_seed(int(seed)) if seed >= 0 else rng.seed()
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=int(num_steps))
+
+    video_info = load_video(video, duration)
+    clip_frames = video_info.clip_frames.unsqueeze(0)
+    sync_frames = video_info.sync_frames.unsqueeze(0)
+
+    seq_cfg.duration = video_info.duration_sec
+    net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+    audio = generate(
+        clip_frames, sync_frames, [prompt],
+        negative_text=[negative_prompt],
+        feature_utils=feature_utils, net=net, fm=fm, rng=rng,
+        cfg_strength=cfg_strength
+    ).float().cpu()[0]
+
+    output_dir.mkdir(exist_ok=True, parents=True)
+    path = output_dir / f"v2a_{datetime.now().strftime('%H%M%S')}.mp4"
+    make_video(video_info, path, audio, sampling_rate=seq_cfg.sampling_rate)
+    gc.collect()
+    return path
+
+@spaces.GPU()
+@torch.inference_mode()
+def image_to_audio(image, prompt, negative_prompt, seed, num_steps, cfg_strength, duration):
+    rng = torch.Generator(device=device)
+    rng.manual_seed(int(seed)) if seed >= 0 else rng.seed()
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=int(num_steps))
+
+    image_info = load_image(image)
+    clip_frames = image_info.clip_frames.unsqueeze(0)
+    sync_frames = image_info.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)
+
+    audio = generate(
+        clip_frames, sync_frames, [prompt],
+        negative_text=[negative_prompt],
+        feature_utils=feature_utils, net=net, fm=fm, rng=rng,
+        cfg_strength=cfg_strength, image_input=True
+    ).float().cpu()[0]
+
+    output_dir.mkdir(exist_ok=True, parents=True)
+    path = output_dir / f"i2a_{datetime.now().strftime('%H%M%S')}.mp4"
+    video_info = VideoInfo.from_image_info(image_info, duration, fps=Fraction(1))
+    make_video(video_info, path, audio, sampling_rate=seq_cfg.sampling_rate)
+    gc.collect()
+    return path
+
+@spaces.GPU()
+@torch.inference_mode()
+def text_to_audio(prompt, negative_prompt, seed, num_steps, cfg_strength, duration):
+    rng = torch.Generator(device=device)
+    rng.manual_seed(int(seed)) if seed >= 0 else rng.seed()
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=int(num_steps))
+
+    seq_cfg.duration = duration
+    net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+    audio = generate(
+        None, None, [prompt],
+        negative_text=[negative_prompt],
+        feature_utils=feature_utils, net=net, fm=fm, rng=rng,
+        cfg_strength=cfg_strength
+    ).float().cpu()[0]
+
+    output_dir.mkdir(exist_ok=True, parents=True)
+    path = output_dir / f"t2a_{datetime.now().strftime('%H%M%S')}.flac"
+    torchaudio.save(path, audio, seq_cfg.sampling_rate)
+    gc.collect()
+    return path
+
 # --- GRADIO UI ---
 
 with gr.Blocks() as demo: