CPU-first startup: load all models/scenes to CPU at module level, GPU transfer at generation time

- startup() now uses device="cpu" for StreamingVideoGenerator and InteractiveSimulator
- Added move_pipeline_to_device() to StreamingVideoGenerator
- Added move_to_device() to _MinimalSVR and InteractiveSimulator
- do_generate() transfers everything to GPU at start, back to CPU in finally
- Warmup deferred to first generation call (CUDA kernel compile on GPU)
- Avoids ZeroGPU time limit: only fast tensor moves happen inside GPU slot

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app.py

@@ -286,6 +286,7 @@ def startup():
         use_ema=USE_EMA,
         seed=SEED,
         enable_taehv=ENABLE_TAEHV,
+        device="cpu",
     )
     video_generator.setup()
     log_gpu("after video generator setup")
@@ -303,7 +304,7 @@ def startup():
         if case_name == "santa_cloth":
             config_overrides["skip_force_fields"] = True
 
-        simulator = InteractiveSimulator(str(case_dir), config_overrides=config_overrides)
+        simulator = InteractiveSimulator(str(case_dir), device="cpu", config_overrides=config_overrides)
         simulator.config["debug"] = False
         log_gpu(f"after simulator init ({case_name})")
 
@@ -350,17 +351,11 @@ def startup():
     log_gpu("after finish_precompute")
 
     # ---- Step 5: Warmup ----
-    first_case = list(cases.keys())[0]
-    _warmup_pipeline(first_case)
-
-    # Release per-case precomputed tensors (i2v_conditional, full_y,
-    # default_text_features) back to CPU so ZeroGPU can reclaim VRAM
-    # between the startup @spaces.GPU slot and the first generation call.
-    # The main model weights (transformer, VAE, text_encoder) stay on GPU;
-    # ZeroGPU preserves the process's GPU slot across calls for efficiency.
-    video_generator.move_case_data_to_device("cpu")
+    # Warmup (CUDA kernel compilation) is deferred to first generation call.
+    print("[5/6] Skipping warmup at CPU-only startup — CUDA kernels compile on first generation.")
+
     torch.cuda.empty_cache()
-    print("[6/6] Startup complete — Gradio server starting.")
+    print("[6/6] CPU-only startup complete — models and scenes ready. GPU transfer at generation time.")
 
 
 # ---------------------------------------------------------------------------
@@ -488,20 +483,17 @@ def do_generate(case_name, prompt, d0, s0, d1, s1, d2, s2):
         _is_generating = True
         _stop_event.clear()
 
-    # Lazy full initialization: load models + build physics scenes + precompute.
-    # Runs only on the first generation; subsequent calls skip this branch.
     if video_generator is None:
-        yield None, "First run: loading models and initializing physics (this takes a minute)..."
-        try:
-            startup()
-        except Exception as e:
-            import traceback; traceback.print_exc()
-            _is_generating = False
-            yield None, f"Initialization error: {e}"
-            return
+        _is_generating = False
+        yield None, "Error: models not initialized. Please reload the Space."
+        return
 
-    # Move precomputed case tensors back to CUDA for this generation session.
+    # Transfer all CPU-resident state to GPU for this generation session.
+    video_generator.move_pipeline_to_device("cuda")
     video_generator.move_case_data_to_device("cuda")
+    for _b in cases.values():
+        if _b.simulator is not None:
+            _b.simulator.move_to_device("cuda")
 
     bundle = cases[case_name]
 
@@ -691,10 +683,12 @@ def do_generate(case_name, prompt, d0, s0, d1, s1, d2, s2):
             render_thread.join(timeout=10)
         if warp_thread is not None:
             warp_thread.join(timeout=10)
-        # Release precomputed case tensors to CPU so ZeroGPU can reclaim
-        # VRAM for other users once this generation session ends.
         if video_generator is not None:
+            video_generator.move_pipeline_to_device("cpu")
             video_generator.move_case_data_to_device("cpu")
+        for _b in cases.values():
+            if _b is not None and _b.simulator is not None:
+                _b.simulator.move_to_device("cpu")
         torch.cuda.empty_cache()
         _is_generating = False
 
@@ -921,7 +915,7 @@ def build_demo():
 # files are already on disk so snapshot_download() is a fast no-op.  By doing
 # this here we avoid holding a ZeroGPU allocation while waiting on downloads.
 _ensure_models_downloaded()
-
+startup()   # Load all models and scenes to CPU at module level
 demo = build_demo()
 
 if __name__ == "__main__":

simulation_engine.py

@@ -306,6 +306,34 @@ class InteractiveSimulator:
     def set_demo_case_handler(self, handler):
         self.demo_case_handler = handler
 
+    def move_to_device(self, device):
+        """Move all renderer/simulation tensors to target device (CPU↔GPU)."""
+        dev = torch.device(device)
+        self.device = dev
+        # Move SVR (PyTorch3D renderer + camera + point clouds)
+        self.svr.move_to_device(dev)
+        # Move mesh data
+        for mesh in self.fg_meshes:
+            for k, v in list(mesh.items()):
+                if isinstance(v, torch.Tensor):
+                    mesh[k] = v.to(dev)
+        # Move foreground point clouds
+        for pc_list in (self.fg_pcs_pt3d, self.fg_pcs_gs):
+            for pc in pc_list:
+                for k, v in list(pc.items()):
+                    if isinstance(v, torch.Tensor):
+                        pc[k] = v.to(dev)
+        # Move per-object transform matrices and initial particles
+        for k in list(self.initial_transform_matrix.keys()):
+            self.initial_transform_matrix[k] = self.initial_transform_matrix[k].to(dev)
+        for k in list(self._init_particles_gpu.keys()):
+            self._init_particles_gpu[k] = self._init_particles_gpu[k].to(dev)
+        # Move obj_info tensors (shared with case_handler by reference)
+        for obj_info in self.all_obj_info:
+            for k, v in list(obj_info.items()):
+                if isinstance(v, torch.Tensor):
+                    obj_info[k] = v.to(dev)
+
     def _load_object_masks(self):
         masks_dir = self.demo_data_path / "fg_masks"
         if not masks_dir.exists():
@@ -615,6 +643,27 @@ class _MinimalSVR:
             compositor=AlphaCompositor(),
         )
 
+    def move_to_device(self, device):
+        """Move all tensors to target device and rebuild renderers."""
+        from pytorch3d.renderer import PerspectiveCameras
+        cam = self.current_camera
+        self.current_camera = PerspectiveCameras(
+            K=cam.K.to(device),
+            R=cam.R.to(device),
+            T=cam.T.to(device),
+            in_ndc=False,
+            image_size=((512, 512),),
+            device=device,
+        )
+        self.bg_points = self.bg_points.to(device)
+        self.bg_points_colors = self.bg_points_colors.to(device)
+        for pc in self.fg_pcs:
+            pc['points'] = pc['points'].to(device)
+            pc['colors'] = pc['colors'].to(device)
+        self.device = device
+        self.cache_bg = None  # stale after device change; recomputed on next render
+        self._build_cached_renderers()
+
     def update_fg_obj_info(self, all_obj_points):
         for idx, pts in enumerate(all_obj_points):
             self.fg_pcs[idx]["points"] = pts.clone()

video_generator.py

@@ -109,13 +109,13 @@ class StreamingVideoGenerator:
         log_gpu("after checkpoint load (bf16, CPU)")
 
         if low_memory:
-            DynamicSwapInstaller.install_model(self.pipeline.text_encoder, device=gpu)
+            DynamicSwapInstaller.install_model(self.pipeline.text_encoder, device=self.device)
         else:
-            self.pipeline.text_encoder.to(device=gpu)
+            self.pipeline.text_encoder.to(device=self.device)
 
-        self.pipeline.generator.to(device=gpu)
-        self.pipeline.vae.to(device=gpu)
-        self.pipeline.encode_vae.to(device=gpu, dtype=torch.bfloat16)
+        self.pipeline.generator.to(device=self.device)
+        self.pipeline.vae.to(device=self.device)
+        self.pipeline.encode_vae.to(device=self.device, dtype=torch.bfloat16)
 
         if self.enable_taehv:
             import os
@@ -138,9 +138,9 @@ class StreamingVideoGenerator:
             self.taehv_decoder.requires_grad_(False)
 
         self.pipeline.processor_dtype = torch.float32
-        self.pipeline.processor_device = gpu
-        self.pipeline.processor_vae = WanVideoVAE().to(device=gpu, dtype=torch.float32)
-        self.pipeline.processor_ienc = WanImageEncoder().to(device=gpu, dtype=torch.float32)
+        self.pipeline.processor_device = self.device
+        self.pipeline.processor_vae = WanVideoVAE().to(device=self.device, dtype=torch.float32)
+        self.pipeline.processor_ienc = WanImageEncoder().to(device=self.device, dtype=torch.float32)
 
         self.pipeline.processor_vae.requires_grad_(False)
         self.pipeline.processor_ienc.requires_grad_(False)
@@ -477,3 +477,17 @@ class StreamingVideoGenerator:
         self.current_start_frame = 0
         self.conditional_dict = None
         self.taehv_cache = None
+
+    def move_pipeline_to_device(self, device: str):
+        """Move all pipeline models to target device (CPU→GPU at generation start, GPU→CPU at end)."""
+        dev = torch.device(device)
+        self.device = dev
+        pipeline = self.pipeline
+        if hasattr(pipeline, 'generator') and pipeline.generator is not None:
+            pipeline.generator.to(device=dev)
+        if hasattr(pipeline, 'vae') and pipeline.vae is not None:
+            pipeline.vae.to(device=dev)
+        if hasattr(pipeline, 'encode_vae') and pipeline.encode_vae is not None:
+            pipeline.encode_vae.to(device=dev)
+        if hasattr(pipeline, 'text_encoder') and pipeline.text_encoder is not None:
+            pipeline.text_encoder.to(device=dev)