Update app.py

app.py

@@ -10,11 +10,13 @@ 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")
 
-# 1. Setup paths so 'mmaudio' package is found
 current_dir = os.path.dirname(os.path.abspath(__file__))
-# Assumes your weights are in a folder named 'weights' in the root
-WEIGHTS_DIR = os.path.join(current_dir, "weights") 
 sys.path.insert(0, os.path.join(current_dir, "MMAudio"))
 
 from mmaudio.eval_utils import (ModelConfig, VideoInfo, all_model_cfg, generate, load_image,
@@ -33,29 +35,33 @@ 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')
 
-# 2. Configure Model Paths (Overriding the auto-downloader)
-model: ModelConfig = all_model_cfg['large_44k_v2']
-# Set these to your local paths manually
-model.model_path = os.path.join(WEIGHTS_DIR, "large_44k_v2.pth")
-model.vae_path = os.path.join(WEIGHTS_DIR, "vggsound_fp32.pth")
-model.synchformer_ckpt = os.path.join(WEIGHTS_DIR, "synchformer_state_dict.pth")
-model.bigvgan_16k_path = os.path.join(WEIGHTS_DIR, "bigvgan_16k_v2.pth")
+# --- WEIGHT SYNCHRONIZATION ---
+def get_weights():
+    log.info(f"Syncing weights from {MY_VAULT_REPO}...")
+    # Adjust filenames below to match your actual vault structure
+    return {
+        "model": hf_hub_download(repo_id=MY_VAULT_REPO, filename="weights/large_44k_v2.pth", token=HF_TOKEN),
+        "vae": hf_hub_download(repo_id=MY_VAULT_REPO, filename="ext_weights/vggsound_fp32.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/bigvgan_16k_v2.pth", token=HF_TOKEN)
+    }
 
-output_dir = Path('./output/gradio')
+weight_paths = get_weights()
 
-def get_model() -> tuple[MMAudio, FeaturesUtils, SequenceConfig]:
-    seq_cfg = model.seq_cfg
+# --- 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"]
 
-    # Load Main Network
+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()
-    if os.path.exists(model.model_path):
-        net.load_weights(torch.load(model.model_path, map_location=device, weights_only=True))
-        log.info(f'Loaded weights from {model.model_path}')
-    else:
-        log.error(f"WEIGHTS NOT FOUND AT: {model.model_path}")
-
-    # Load Feature Utils (Synchformer, VAE, etc.)
+    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,
@@ -63,18 +69,17 @@ def get_model() -> tuple[MMAudio, FeaturesUtils, SequenceConfig]:
         mode=model.mode,
         bigvgan_vocoder_ckpt=model.bigvgan_16k_path,
         need_vae_encoder=False
-    )
-    feature_utils = feature_utils.to(device, dtype).eval()
+    ).to(device, dtype).eval()
 
     return net, feature_utils, seq_cfg
 
-# Initialize global model components
-net, feature_utils, seq_cfg = get_model()
+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):
-    # Rest of your existing logic remains the same...
     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))
@@ -94,36 +99,98 @@ def video_to_audio(video, prompt, negative_prompt, seed, num_steps, cfg_strength
     ).float().cpu()[0]
 
     output_dir.mkdir(exist_ok=True, parents=True)
-    path = output_dir / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
+    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
 
-# ... [Keep your image_to_audio and text_to_audio functions here] ...
+@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 ---
+
+common_inputs = [
+    gr.Text(label="Prompt"),
+    gr.Text(label="Negative Prompt", value="music"),
+    gr.Number(label="Seed (-1 = random)", value=-1),
+    gr.Slider(label="Num Steps", minimum=1, maximum=100, value=25, step=1),
+    gr.Slider(label="Guidance Strength", minimum=1, maximum=15, value=4.5, step=0.1),
+    gr.Number(label="Duration (sec)", value=8),
+]
 
-# UI Section
 video_to_audio_tab = gr.Interface(
     fn=video_to_audio,
-    inputs=[
-        gr.Video(label="Video"),
-        gr.Text(label="Prompt"),
-        gr.Text(label="Negative Prompt", value="music"),
-        gr.Number(label="Seed (-1 = random)", value=-1),
-        gr.Slider(label="Num Steps", minimum=1, maximum=100, value=25, step=1),
-        gr.Slider(label="Guidance Strength", minimum=1, maximum=15, value=4.5, step=0.1),
-        gr.Number(label="Max Duration (sec)", value=8),
-    ],
-    outputs=gr.Video(label="Generated Video"),
+    inputs=[gr.Video(label="Input Video")] + common_inputs,
+    outputs=gr.Video(label="Generated Video with Audio"),
     title="Video-to-Audio"
 )
 
-# ... [Keep your other tabs] ...
+image_to_audio_tab = gr.Interface(
+    fn=image_to_audio,
+    inputs=[gr.Image(label="Input Image", type="filepath")] + common_inputs,
+    outputs=gr.Video(label="Static Video with Audio"),
+    title="Image-to-Audio"
+)
+
+text_to_audio_tab = gr.Interface(
+    fn=text_to_audio,
+    inputs=common_inputs[:5] + [gr.Number(label="Duration (sec)", value=8)],
+    outputs=gr.Audio(label="Generated Audio"),
+    title="Text-to-Audio"
+)
 
 if __name__ == "__main__":
     output_dir.mkdir(exist_ok=True, parents=True)
     interface = gr.TabbedInterface(
-        [video_to_audio_tab], # Add back text_to_audio_tab, image_to_audio_tab
-        ["Video-to-Audio"]
+        [video_to_audio_tab, image_to_audio_tab, text_to_audio_tab],
+        ["Video-to-Audio", "Image-to-Audio", "Text-to-Audio"]
     )
     interface.launch(server_name="0.0.0.0", server_port=7860)