Update app.py

app.py

@@ -124,6 +124,20 @@ print("Pipeline ready!")
 print("=" * 80)
 
 
+def log_memory(tag: str):
+    """Log GPU memory usage at a given point."""
+    if torch.cuda.is_available():
+        allocated = torch.cuda.memory_allocated() / 1024**3
+        reserved = torch.cuda.memory_reserved() / 1024**3
+        peak = torch.cuda.max_memory_allocated() / 1024**3
+        free, total = torch.cuda.mem_get_info()
+        free_gb = free / 1024**3
+        total_gb = total / 1024**3
+        print(f"[VRAM {tag}] allocated={allocated:.2f}GB reserved={reserved:.2f}GB peak={peak:.2f}GB free={free_gb:.2f}GB total={total_gb:.2f}GB")
+    else:
+        print(f"[VRAM {tag}] CUDA not available")
+
+
 @spaces.GPU(duration=120, size="xlarge")
 def generate_video(
     input_image,
@@ -138,6 +152,9 @@ def generate_video(
 ):
     """Generate a video based on the given parameters."""
     try:
+        torch.cuda.reset_peak_memory_stats()
+        log_memory("start")
+
         current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
         frame_rate = DEFAULT_FRAME_RATE
@@ -145,6 +162,8 @@ def generate_video(
         # 8k+1 format
         num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1
 
+        print(f"Generating: {height}x{width}, {num_frames} frames ({duration}s), seed={current_seed}")
+
         # Handle image input
         images = []
         temp_image_path = None
@@ -189,6 +208,7 @@ def generate_video(
             if audio_context is not None:
                 audio_context = audio_context.to("cuda")
             print("Embeddings loaded successfully")
+            log_memory("after embeddings loaded")
         except Exception as e:
             raise RuntimeError(
                 f"Failed to get embeddings from text encoder space: {e}\n"
@@ -209,6 +229,8 @@ def generate_video(
             tiling_config = TilingConfig.default()
             video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
 
+            log_memory("before pipeline call")
+
             video, audio = pipeline(
                 prompt=prompt,
                 seed=current_seed,
@@ -221,6 +243,8 @@ def generate_video(
                 enhance_prompt=False,  # Already enhanced by text encoder space
             )
 
+            log_memory("after pipeline call")
+
             output_path = tempfile.mktemp(suffix=".mp4")
             encode_video(
                 video=video,
@@ -230,6 +254,8 @@ def generate_video(
                 video_chunks_number=video_chunks_number,
             )
 
+            log_memory("after encode_video")
+
             return str(output_path), current_seed
         finally:
             # Restore original encode_prompts
@@ -237,6 +263,7 @@ def generate_video(
 
     except Exception as e:
         import traceback
+        log_memory("on error")
         error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
         print(error_msg)
         return None, current_seed