fix: load only AudioLDM2 VAE+vocoder subcomponents instead of full pipeline to prevent GPU OOM on long videos

app.py

@@ -313,7 +313,8 @@ def generate_taro(video_file, seed_val, cfg_scale, num_steps, mode,
     from TARO.onset_util import VideoOnsetNet, extract_onset
     from TARO.models    import MMDiT
     from TARO.samplers  import euler_sampler, euler_maruyama_sampler
-    from diffusers      import AudioLDM2Pipeline
+    from diffusers      import AutoencoderKL
+    from transformers   import SpeechT5HifiGan
 
     # -- Load CAVP encoder (uses checkpoint from our HF repo) --
     extract_cavp = Extract_CAVP_Features(
@@ -343,11 +344,10 @@ def generate_taro(video_file, seed_val, cfg_scale, num_steps, mode,
     model.load_state_dict(torch.load(taro_ckpt_path, map_location=device, weights_only=False)["ema"])
     model.eval().to(weight_dtype)
 
-    # -- Load AudioLDM2 VAE + vocoder (decoder pipeline only) --
-    # TARO uses AudioLDM2's VAE and vocoder for decoding; no encoder needed at inference
-    audioldm2 = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2")
-    vae        = audioldm2.vae.to(device).eval()
-    vocoder    = audioldm2.vocoder.to(device)
+    # -- Load AudioLDM2 VAE + vocoder only (saves ~3-4 GB vs loading the full pipeline) --
+    # TARO only needs VAE and vocoder for decoding; the text encoder and UNet are never used.
+    vae     = AutoencoderKL.from_pretrained("cvssp/audioldm2", subfolder="vae").to(device).eval()
+    vocoder = SpeechT5HifiGan.from_pretrained("cvssp/audioldm2", subfolder="vocoder").to(device)
     latents_scale = torch.tensor([0.18215] * 8).view(1, 8, 1, 1).to(device)
 
     # -- Prepare silent video (shared across all samples) --
@@ -873,7 +873,8 @@ def regen_taro_segment(video_file, seg_idx, seg_meta_json,
     from TARO.onset_util import VideoOnsetNet, extract_onset
     from TARO.models    import MMDiT
     from TARO.samplers  import euler_sampler, euler_maruyama_sampler
-    from diffusers      import AudioLDM2Pipeline
+    from diffusers      import AutoencoderKL
+    from transformers   import SpeechT5HifiGan
 
     silent_video  = meta["silent_video"]
     tmp_dir       = tempfile.mkdtemp()
@@ -891,9 +892,8 @@ def regen_taro_segment(video_file, seg_idx, seg_meta_json,
     model_net     = MMDiT(adm_in_channels=120, z_dims=[768], encoder_depth=4).to(device)
     model_net.load_state_dict(torch.load(taro_ckpt_path, map_location=device, weights_only=False)["ema"])
     model_net.eval().to(weight_dtype)
-    audioldm2     = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2")
-    vae           = audioldm2.vae.to(device).eval()
-    vocoder       = audioldm2.vocoder.to(device)
+    vae           = AutoencoderKL.from_pretrained("cvssp/audioldm2", subfolder="vae").to(device).eval()
+    vocoder       = SpeechT5HifiGan.from_pretrained("cvssp/audioldm2", subfolder="vocoder").to(device)
     latents_scale = torch.tensor([0.18215] * 8).view(1, 8, 1, 1).to(device)
 
     cavp_feats    = extract_cavp(silent_video, tmp_path=tmp_dir)