Upload VideoEngine.py

VideoEngine.py

@@ -134,14 +134,14 @@ class VideoEngine:
         if self.device != "cuda":
             return
 
-        print("→ [4/7] Enabling TF32 acceleration...")
+        print("→ [5/8] Enabling TF32 acceleration...")
         torch.backends.cuda.matmul.allow_tf32 = True
         torch.backends.cudnn.allow_tf32 = True
         print("✓ TF32 enabled (20-30% speedup on A100/H100)")
 
     def _enable_xformers(self):
         """Enable xFormers memory efficient attention."""
-        print("→ [5/7] Enabling xFormers attention...")
+        print("→ [6/8] Enabling xFormers attention...")
 
         if not self._check_xformers_available():
             print("⚠ xFormers not available, using standard attention")
@@ -157,10 +157,10 @@ class VideoEngine:
     def _apply_aot_compilation(self):
         """Apply AOTI pre-compiled blocks (HF Spaces deployment only)."""
         if not self.is_spaces:
-            print("→ [6/7] Skipping AOTI (testing mode, not needed)")
+            print("→ [7/8] Skipping AOTI (testing mode, not needed)")
             return
 
-        print("→ [6/7] Loading AOTI pre-compiled blocks...")
+        print("→ [7/8] Loading AOTI pre-compiled blocks...")
         try:
             # Lazy import to avoid CUDA initialization at module load
             from spaces.zero.torch.aoti import aoti_blocks_load
@@ -178,7 +178,7 @@ class VideoEngine:
             aoti_blocks_load(self.pipeline.transformer_2, 'zerogpu-aoti/Wan2', variant=variant)
             print(f"✓ AOTI blocks loaded (variant: {variant}, 60-70% speedup)")
         except ImportError:
-            print("⚠ [6/7] Skipping AOTI (spaces.zero.torch.aoti not available)")
+            print("⚠ [7/8] Skipping AOTI (spaces.zero.torch.aoti not available)")
         except Exception as e:
             print(f"⚠ AOTI load failed (falling back to standard inference): {e}")
             print("  This is not critical, speed will be slightly slower")
@@ -207,50 +207,41 @@ class VideoEngine:
             print("Loading Wan2.2 I2V Engine with Optimizations")
             print("=" * 60)
 
-            # Stage 1: Load base pipeline
-            print("→ [1/7] Loading base pipeline...")
+            # Stage 1: Load base pipeline (keep on CPU first to save VRAM)
+            print("→ [1/8] Loading base pipeline...")
             self.pipeline = WanImageToVideoPipeline.from_pretrained(
                 self.MODEL_ID,
                 transformer=WanTransformer3DModel.from_pretrained(
                     self.TRANSFORMER_REPO,
                     subfolder='transformer',
                     torch_dtype=torch.bfloat16,
-                    device_map='cuda',
+                    low_cpu_mem_usage=True,
                 ),
                 transformer_2=WanTransformer3DModel.from_pretrained(
                     self.TRANSFORMER_REPO,
                     subfolder='transformer_2',
                     torch_dtype=torch.bfloat16,
-                    device_map='cuda',
+                    low_cpu_mem_usage=True,
                 ),
                 torch_dtype=torch.bfloat16,
-            ).to('cuda')
+            )
 
-            # Stage 2: Load Lightning LoRA
-            print("→ [2/7] Loading Lightning LoRA adapters...")
+            # Stage 2: Load Lightning LoRA on CPU first, then fuse
+            print("→ [2/8] Loading Lightning LoRA adapters...")
             self.pipeline.load_lora_weights(self.LORA_REPO, weight_name=self.LORA_WEIGHT, adapter_name="lightx2v")
+            gc.collect()
             kwargs_lora = {"load_into_transformer_2": True}
             self.pipeline.load_lora_weights(self.LORA_REPO, weight_name=self.LORA_WEIGHT, adapter_name="lightx2v_2", **kwargs_lora)
+            gc.collect()
             self.pipeline.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1., 1.])
             self.pipeline.fuse_lora(adapter_names=["lightx2v"], lora_scale=3., components=["transformer"])
             self.pipeline.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1., components=["transformer_2"])
             self.pipeline.unload_lora_weights()
+            gc.collect()
             print("✓ Lightning LoRA loaded and fused")
 
-            # Stage 3: Apply quantization (both environments - critical for memory)
-            self._apply_quantization()
-
-            # Stage 4: Enable TF32 acceleration
-            self._enable_tf32()
-
-            # Stage 5: Enable xFormers attention
-            self._enable_xformers()
-
-            # Stage 6: Apply AOTI (deployment only)
-            self._apply_aot_compilation()
-
-            # Stage 7: Enable VAE optimizations (critical for ZeroGPU memory)
-            print("→ [7/8] Enabling VAE memory optimizations...")
+            # Stage 3: Enable VAE optimizations BEFORE moving to GPU
+            print("→ [3/8] Enabling VAE memory optimizations...")
             try:
                 if hasattr(self.pipeline, 'enable_vae_tiling'):
                     self.pipeline.enable_vae_tiling()
@@ -261,6 +252,23 @@ class VideoEngine:
             except Exception as e:
                 print(f"⚠ VAE optimization warning: {e}")
 
+            # Stage 4: Enable model CPU offload for ZeroGPU memory efficiency
+            print("→ [4/8] Enabling model CPU offload...")
+            self.pipeline.enable_model_cpu_offload()
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            print("✓ Model CPU offload enabled (dynamic GPU/CPU memory management)")
+
+            # Stage 5: Enable TF32 acceleration
+            self._enable_tf32()
+
+            # Stage 6: Enable xFormers attention
+            self._enable_xformers()
+
+            # Stage 7: Apply AOTI (deployment only, usually skipped)
+            self._apply_aot_compilation()
+
             # Stage 8: Final setup
             print("→ [8/8] Finalizing setup...")
             self.is_loaded = True