Upload folder using huggingface_hub

app.py

@@ -1,5 +1,5 @@
 import os
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 import spaces
 import sys
 import time
@@ -447,7 +447,8 @@ def infer_chunk(
 
 def decode_video(pred_latent):
     """Decode latent tensor to video frames."""
-    latents = pred_latent.to(vae.dtype)
+    vae_device = next(vae.parameters()).device
+    latents = pred_latent.to(vae_device).to(vae.dtype)
     latents_mean = (
         torch.tensor(vae.config.latents_mean)
         .view(1, vae.config.z_dim, 1, 1, 1)
@@ -546,13 +547,17 @@ def generate(
 
     pred_latent = torch.cat(pred_latent_lst, dim=2)
 
-    # Clean up caches
+    # Clean up caches and free VRAM before VAE decode
     transformer.clear_cache("pos")
     streaming_vae.clear_cache()
     torch.cuda.empty_cache()
 
-    # Decode video
-    video = decode_video(pred_latent)
+    # Decode video - move pred_latent to CPU, then decode frame-by-frame
+    pred_latent_cpu = pred_latent.cpu()
+    del pred_latent
+    torch.cuda.empty_cache()
+
+    video = decode_video(pred_latent_cpu)
 
     # Save to temp file
     from diffusers.utils import export_to_video