import gc import logging from argparse import ArgumentParser from datetime import datetime from fractions import Fraction from pathlib import Path import gradio as gr import soundfile as sf import spaces import torch from huggingface_hub import hf_hub_download from safetensors.torch import load_file 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 logging.getLogger("httpx").setLevel(logging.WARNING) logging.getLogger("requests").setLevel(logging.WARNING) logging.getLogger("urllib3").setLevel(logging.WARNING) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True log = logging.getLogger() device = 'cuda' if torch.cuda.is_available() else 'cpu' dtype = torch.float32 # Every weight comes from our own namespace; the CLIP encoder and the 44kHz vocoder are repointed # inside the vendored `mmaudio` package (features_utils.py / ext/autoencoder.py). EXT_REPO = 'thornmaze/mma-ext' CORE_REPO = 'thornmaze/mma-core-fp16' # Opaque on purpose: this file is served publicly with the Space, so no name here may carry purpose. CORE_FILE = 'core_44k_fp16.safetensors' # Architecture is large_44k (v1); only the sequence/mode config is taken from the v2 entry. MY_MODEL_NAME = 'large_44k' vae_path = hf_hub_download(repo_id=EXT_REPO, subfolder='ext_weights', filename='v1-44.pth') synchformer_path = hf_hub_download(repo_id=EXT_REPO, subfolder='ext_weights', filename='synchformer_state_dict.pth') model_cfg_for_params: ModelConfig = all_model_cfg['large_44k_v2'] output_dir = Path('./output/gradio') setup_eval_logging() def get_model() -> tuple[MMAudio, FeaturesUtils, SequenceConfig]: seq_cfg = model_cfg_for_params.seq_cfg net: MMAudio = get_my_mmaudio(MY_MODEL_NAME).to(device, dtype).eval() weights_path = hf_hub_download(repo_id=CORE_REPO, filename=CORE_FILE) state_dict = load_file(weights_path) missing, unexpected = net.load_state_dict(state_dict, strict=False) if missing: log.warning(f"Missing keys: {missing[:8]}{'...' if len(missing) > 8 else ''}") if unexpected: log.warning(f"Unexpected keys: {unexpected[:8]}{'...' if len(unexpected) > 8 else ''}") if dtype == torch.float16: net.half() net.to(device).eval() # bigvgan_vocoder_ckpt=None -> the 44kHz BigVGANv2 is pulled from our mirror by the vendored # autoencoder; need_vae_encoder=False because inference only decodes. feature_utils = FeaturesUtils(tod_vae_ckpt=vae_path, synchformer_ckpt=synchformer_path, enable_conditions=True, mode=model_cfg_for_params.mode, bigvgan_vocoder_ckpt=None, need_vae_encoder=False) feature_utils = feature_utils.to(device, dtype).eval() return net, feature_utils, seq_cfg net, feature_utils, seq_cfg = get_model() def _rng(seed: int) -> torch.Generator: rng = torch.Generator(device=device) if seed >= 0: rng.manual_seed(seed) else: rng.seed() return rng def _stamp() -> str: return datetime.now().strftime('%Y%m%d_%H%M%S_%f') def _audio_for_video(video, prompt: str, negative_prompt: str, seed: int, num_steps: int, cfg_strength: float, duration: float): """Shared core of both video paths: decoded video -> generated waveform.""" fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=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) audios = generate(clip_frames, sync_frames, [prompt], negative_text=[negative_prompt], feature_utils=feature_utils, net=net, fm=fm, rng=_rng(seed), cfg_strength=cfg_strength) return video_info, audios.float().cpu()[0] @spaces.GPU(duration=90) @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): video_info, audio = _audio_for_video(video, prompt, negative_prompt, seed, num_steps, cfg_strength, duration) output_dir.mkdir(exist_ok=True, parents=True) video_save_path = output_dir / f'{_stamp()}.mp4' make_video(video_info, video_save_path, audio, sampling_rate=seq_cfg.sampling_rate) gc.collect() return video_save_path @spaces.GPU(duration=90) @torch.inference_mode() def video_to_track(video: gr.Video, prompt: str, negative_prompt: str, seed: int, num_steps: int, cfg_strength: float, duration: float): """Same generation, but returns the bare audio track instead of a re-encoded mp4. `make_video` re-encodes every frame (measured: a 1.4 MB source came back at 6.5 MB with a second lossy pass). The caller muxes this track onto its own untouched video stream. """ _, audio = _audio_for_video(video, prompt, negative_prompt, seed, num_steps, cfg_strength, duration) output_dir.mkdir(exist_ok=True, parents=True) audio_save_path = output_dir / f'{_stamp()}.flac' # torchaudio.save routes through torchcodec, which is not installed on Spaces sf.write(audio_save_path, audio.transpose(0, 1).numpy(), seq_cfg.sampling_rate) gc.collect() return audio_save_path @spaces.GPU(duration=90) @torch.inference_mode() def image_to_audio(image: gr.Image, prompt: str, negative_prompt: str, seed: int, num_steps: int, cfg_strength: float, duration: float): fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=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) audios = generate(clip_frames, sync_frames, [prompt], negative_text=[negative_prompt], feature_utils=feature_utils, net=net, fm=fm, rng=_rng(seed), cfg_strength=cfg_strength, image_input=True) audio = audios.float().cpu()[0] output_dir.mkdir(exist_ok=True, parents=True) video_save_path = output_dir / f'{_stamp()}.mp4' video_info = VideoInfo.from_image_info(image_info, duration, fps=Fraction(1)) make_video(video_info, video_save_path, audio, sampling_rate=seq_cfg.sampling_rate) gc.collect() return video_save_path @spaces.GPU(duration=45) @torch.inference_mode() def text_to_audio(prompt: str, negative_prompt: str, seed: int, num_steps: int, cfg_strength: float, duration: float): fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=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) audios = generate(None, None, [prompt], negative_text=[negative_prompt], feature_utils=feature_utils, net=net, fm=fm, rng=_rng(seed), cfg_strength=cfg_strength) audio = audios.float().cpu()[0] output_dir.mkdir(exist_ok=True, parents=True) audio_save_path = output_dir / f'{_stamp()}.flac' sf.write(audio_save_path, audio.transpose(0, 1).numpy(), seq_cfg.sampling_rate) gc.collect() return audio_save_path def _params(with_negative_default: bool) -> list: return [ gr.Text(label='Prompt'), gr.Text(label='Negative prompt', value='music' if with_negative_default else ''), gr.Number(label='Seed (-1: random)', value=-1, precision=0, minimum=-1), gr.Number(label='Num steps', value=25, precision=0, minimum=1), gr.Number(label='Guidance Strength', value=4.5, minimum=1), gr.Number(label='Duration (sec)', value=8, minimum=1), ] # api_name is set explicitly on every tab: the default names endpoints by tab ORDER # (predict, predict_1, ...), so adding or reordering a tab silently re-points every caller. video_to_audio_tab = gr.Interface( fn=video_to_audio, api_name='v2a', inputs=[gr.Video(), *_params(True)], outputs='playable_video', cache_examples=False, title='Video to Audio', description='Resolutions above 384 px on the shorter side cost time without improving output.', ) video_to_track_tab = gr.Interface( fn=video_to_track, api_name='v2track', inputs=[gr.Video(), *_params(True)], outputs='audio', cache_examples=False, title='Video to Audio track', description='Returns the generated track only — the source video is never re-encoded.', ) text_to_audio_tab = gr.Interface( fn=text_to_audio, api_name='t2a', inputs=_params(False), outputs='audio', cache_examples=False, title='Text to Audio', ) image_to_audio_tab = gr.Interface( fn=image_to_audio, api_name='i2a', inputs=[gr.Image(type='filepath'), *_params(False)], outputs='playable_video', cache_examples=False, title='Image to Audio (experimental)', ) app = gr.TabbedInterface( [video_to_audio_tab, video_to_track_tab, text_to_audio_tab, image_to_audio_tab], ['Video-to-Audio', 'Video-to-Track', 'Text-to-Audio', 'Image-to-Audio'], ) if __name__ == "__main__": parser = ArgumentParser() parser.add_argument('--port', type=int, default=7860) parser.add_argument('--share', action='store_true', help='Create a public link') args = parser.parse_args() app.launch(server_name="0.0.0.0", server_port=args.port, share=args.share, allowed_paths=[output_dir])