Update app.py

app.py

@@ -169,6 +169,27 @@ ASPECT_RATIOS = {
     }
 }
 
+def get_vae_cache_for_aspect_ratio(aspect_ratio, device, dtype):
+    """
+    Create VAE cache with appropriate dimensions for the given aspect ratio.
+    """
+    ar_config = ASPECT_RATIOS[aspect_ratio]
+    latent_h = ar_config["latent_h"]
+    latent_w = ar_config["latent_w"]
+    
+    # Create new cache tensors with correct dimensions
+    # Based on ZERO_VAE_CACHE structure but adjusted for aspect ratio
+    cache = []
+    
+    # The cache dimensions need to match the latent dimensions
+    # These are placeholder tensors that will be updated during generation
+    cache.append(torch.zeros(1, 512, latent_h // 8, latent_w // 8, device=device, dtype=dtype))  # 8x downsampled
+    cache.append(torch.zeros(1, 512, latent_h // 4, latent_w // 4, device=device, dtype=dtype))  # 4x downsampled
+    cache.append(torch.zeros(1, 256, latent_h // 2, latent_w // 2, device=device, dtype=dtype))  # 2x downsampled
+    cache.append(torch.zeros(1, 128, latent_h, latent_w, device=device, dtype=dtype))            # 1x (same as latent)
+    
+    return cache
+
 def frames_to_ts_file(frames, filepath, fps = 15):
     """
     Convert frames directly to .ts file using PyAV.
@@ -360,7 +381,8 @@ def video_generation_handler_streaming(prompt, seed=42, fps=15, aspect_ratio="16
     
     vae_cache, latents_cache = None, None
     if not APP_STATE["current_use_taehv"] and not args.trt:
-        vae_cache = [c.to(device=gpu, dtype=torch.float16) for c in ZERO_VAE_CACHE]
+        # Create VAE cache with correct dimensions for the aspect ratio
+        vae_cache = get_vae_cache_for_aspect_ratio(aspect_ratio, gpu, torch.float16)
 
     num_blocks = 7
     current_start_frame = 0