Add app.py with Gradio interface for TARO inference

app.py

@@ -0,0 +1,95 @@
+import os
+import torch
+import numpy as np
+import random
+import soundfile as sf
+import ffmpeg
+import tempfile
+import spaces
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+REPO_ID = "JackIsNotInTheBox/Taro_checkpoints"
+CACHE_DIR = "/tmp/taro_ckpts"
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+print("Downloading checkpoints...")
+cavp_ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="cavp_epoch66.ckpt", cache_dir=CACHE_DIR)
+onset_ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="onset_model.ckpt", cache_dir=CACHE_DIR)
+taro_ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="taro_ckpt.pt", cache_dir=CACHE_DIR)
+print("Checkpoints downloaded.")
+
+def set_global_seed(seed):
+    np.random.seed(seed % (2**32))
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+@spaces.GPU(duration=300)
+def generate_audio(video_file, seed_val, cfg_scale, num_steps, mode):
+    set_global_seed(int(seed_val))
+    torch.set_grad_enabled(False)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    weight_dtype = torch.bfloat16
+    from cavp_util import Extract_CAVP_Features
+    from onset_util import VideoOnsetNet, extract_onset
+    from models import MMDiT
+    from samplers import euler_sampler, euler_maruyama_sampler
+    from diffusers import AudioLDM2Pipeline
+    extract_cavp = Extract_CAVP_Features(device=device, config_path="./cavp/cavp.yaml", ckpt_path=cavp_ckpt_path)
+    state_dict = torch.load(onset_ckpt_path, map_location=device)["state_dict"]
+    new_state_dict = {}
+    for key, value in state_dict.items():
+        if "model.net.model" in key:
+            new_key = key.replace("model.net.model", "net.model")
+        elif "model.fc." in key:
+            new_key = key.replace("model.fc", "fc")
+        else:
+            new_key = key
+        new_state_dict[new_key] = value
+    onset_model = VideoOnsetNet(False).to(device)
+    onset_model.load_state_dict(new_state_dict)
+    onset_model.eval()
+    model = MMDiT(adm_in_channels=120, z_dims=[768], encoder_depth=4).to(device)
+    ckpt = torch.load(taro_ckpt_path, map_location=device)["ema"]
+    model.load_state_dict(ckpt)
+    model.eval()
+    model.to(weight_dtype)
+    model_audioldm = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2")
+    vae = model_audioldm.vae.to(device)
+    vae.eval()
+    vocoder = model_audioldm.vocoder.to(device)
+    tmp_dir = tempfile.mkdtemp()
+    cavp_feats = extract_cavp(video_file, tmp_path=tmp_dir)
+    onset_feats = extract_onset(video_file, onset_model, tmp_path=tmp_dir, device=device)
+    sr = 16000
+    truncate = 131072
+    fps = 4
+    truncate_frame = int(fps * truncate / sr)
+    truncate_onset = 120
+    latents_scale = torch.tensor([0.18215]*8).view(1, 8, 1, 1).to(device)
+    video_feats = torch.from_numpy(cavp_feats[:truncate_frame]).unsqueeze(0).to(device).to(weight_dtype)
+    onset_feats_t = torch.from_numpy(onset_feats[:truncate_onset]).unsqueeze(0).to(device).to(weight_dtype)
+    z = torch.randn(len(video_feats), model.in_channels, 204, 16, device=device).to(weight_dtype)
+    sampling_kwargs = dict(model=model, latents=z, y=onset_feats_t, context=video_feats, num_steps=int(num_steps), heun=False, cfg_scale=float(cfg_scale), guidance_low=0.0, guidance_high=0.7, path_type="linear")
+    with torch.no_grad():
+        if mode == "sde":
+            samples = euler_maruyama_sampler(**sampling_kwargs)
+        else:
+            samples = euler_sampler(**sampling_kwargs)
+    samples = vae.decode(samples / latents_scale).sample
+    wav_samples = vocoder(samples.squeeze()).detach().cpu().numpy()
+    audio_path = os.path.join(tmp_dir, "output.wav")
+    sf.write(audio_path, wav_samples, sr)
+    duration = truncate / sr
+    trimmed_video = os.path.join(tmp_dir, "trimmed.mp4")
+    output_video = os.path.join(tmp_dir, "output.mp4")
+    ffmpeg.input(video_file, ss=0, t=duration).output(trimmed_video, vcodec="libx264", an=None).run(overwrite_output=True, quiet=True)
+    input_v = ffmpeg.input(trimmed_video)
+    input_a = ffmpeg.input(audio_path)
+    ffmpeg.output(input_v, input_a, output_video, vcodec="libx264", acodec="aac", strict="experimental").run(overwrite_output=True, quiet=True)
+    return output_video, audio_path
+
+demo = gr.Interface(fn=generate_audio, inputs=[gr.Video(label="Input Video"), gr.Number(label="Seed", value=0, precision=0), gr.Slider(label="CFG Scale", minimum=1, maximum=15, value=8, step=0.5), gr.Slider(label="Sampling Steps", minimum=10, maximum=50, value=25, step=1), gr.Radio(label="Sampling Mode", choices=["sde", "ode"], value="sde")], outputs=[gr.Video(label="Output Video with Audio"), gr.Audio(label="Generated Audio")], title="TARO: Video-to-Audio Synthesis (ICCV 2025)", description="Upload a video and generate synchronized audio using TARO.")
+demo.queue().launch()