app.py

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