appearance fixes

app2.py

@@ -946,12 +946,14 @@ def run_rollout(video_path: str, gen_seconds: float, l1_steps: int, l2_steps: in
             traj_path = job / "trajectory.npy"
             np.save(traj_path, trajectory_to_model_frame(trajectory))
 
-        # First request in this worker moves the model to CUDA and compiles it;
-        # every later request just reuses that already-ready model (see
+        # First request in this worker moves the model to CUDA, then compiles it only if
+        # the pre-compiled artifacts below can be loaded -- otherwise it runs eagerly
+        # rather than paying for a from-scratch compile live in a request. Every later
+        # request just reuses that already-ready (compiled or eager) model (see
         # orbis2_app_engine.py).
-        print("[load] ensure_ready: moving model to GPU / compiling if needed…")
-        ENGINE.ensure_ready(device="cuda", compile=False,
-                             compile_artifacts=str(EXP_DIR / "compile_cache.pkl"))
+        print("[load] ensure_ready: moving model to GPU, compiling only if cached artifacts load…")
+        ENGINE.ensure_ready(device="cuda", compile=True,
+                             compile_artifacts=str(EXP_DIR / "checkpoints" / "compile_cache.pkl"))
 
         say(0.10, "🚗 Rolling out…")
 

orbis2_app_engine.py

@@ -118,8 +118,6 @@ class RolloutEngine:
         self.model = None
         self._ready = False
         self._lock = threading.Lock()
-        self._compile_artifacts = None
-        self._artifacts_saved = False
 
     def load(self):
         """CPU-only one-time setup: read the config, build the model, load the
@@ -178,12 +176,15 @@ class RolloutEngine:
 
     def ensure_ready(self, device="cuda", compile=True, compile_mode="reduce-overhead",
                       compile_artifacts=None, speed_scale=1.0, yaw_rate_scale=1.0):
-        """Move the model to `device` and wrap it with torch.compile, exactly once.
+        """Move the model to `device` and, only if a matching pre-compiled artifacts
+        file is available, wrap it with torch.compile and warm-start from that cache.
+        Otherwise runs eagerly -- this never compiles from scratch live in a request,
+        since tracing/codegen from nothing is far too slow to pay on a real request.
         Safe to call on every request: after the first call this is a no-op, because
-        the model lives in this persistent process and keeps its GPU/compiled state
-        between calls."""
+        the model lives in this persistent process and keeps its GPU/compiled-or-eager
+        state between calls."""
         if self._ready:
-            print("[compile] model already on-device and compiled from a prior request -- reusing it as-is.")
+            print("[compile] model already on-device from a prior request -- reusing it as-is.")
             return self.model
 
         with self._lock:
@@ -192,7 +193,25 @@ class RolloutEngine:
 
             model = self.model.to(device)
 
-            if compile:
+            use_compile = False
+            if compile and compile_artifacts and os.path.exists(compile_artifacts):
+                try:
+                    with open(compile_artifacts, "rb") as f:
+                        torch.compiler.load_cache_artifacts(f.read())
+                    use_compile = True
+                    print(f"[compile] loaded cached compile artifacts from {compile_artifacts!r} "
+                          "-- compiling with these instead of tracing from scratch.")
+                    logger.info(f"Loaded compile artifacts from {compile_artifacts!r}")
+                except Exception as e:
+                    print(f"[compile] failed to load compile artifacts from {compile_artifacts!r} "
+                          f"({e!r}) -- running eagerly, no torch.compile.")
+            elif compile and compile_artifacts:
+                print(f"[compile] no compile artifacts found at {compile_artifacts!r} "
+                      "-- running eagerly, no torch.compile (will not compile from scratch).")
+            elif compile:
+                print("[compile] no compile_artifacts path given -- running eagerly, no torch.compile.")
+
+            if use_compile:
                 def _maybe_compile(module, attr):
                     net = getattr(module, attr, None)
                     if net is not None:
@@ -205,29 +224,12 @@ class RolloutEngine:
                 if l2_predictor is not None:
                     _maybe_compile(l2_predictor, "ema_vit")
 
-                if compile_artifacts:
-                    if os.path.exists(compile_artifacts):
-                        with open(compile_artifacts, "rb") as f:
-                            torch.compiler.load_cache_artifacts(f.read())
-                        print(f"[compile] loaded cached compile artifacts from {compile_artifacts!r} "
-                              "-- first call reuses these instead of tracing from scratch.")
-                        logger.info(f"Loaded compile artifacts from {compile_artifacts!r}")
-                    else:
-                        print(f"[compile] no compile artifacts found at {compile_artifacts!r} "
-                              "-- compiling from scratch; will save artifacts after the first rollout.")
-                        logger.info(
-                            f"Compile artifacts not found at {compile_artifacts!r}; "
-                            "will save after the first rollout."
-                        )
-                else:
-                    print("[compile] no compile_artifacts path given -- compiling from scratch, nothing cached to disk.")
-                self._compile_artifacts = compile_artifacts
-
             maybe_apply_condition_preprocessor_scales(model, speed_scale, yaw_rate_scale)
 
             self.model = model
             self._ready = True
-            logger.info(f"Model moved to {device} and ready -- subsequent requests reuse this instance.")
+            logger.info(f"Model moved to {device} and ready "
+                        f"({'compiled' if use_compile else 'eager'}) -- subsequent requests reuse this instance.")
 
         return self.model
 
@@ -368,14 +370,6 @@ class RolloutEngine:
         if device.startswith("cuda"):
             logger.info(f"Max memory: {torch.cuda.max_memory_allocated() / 1024**3:.02f} GB")
 
-        if self._compile_artifacts and not self._artifacts_saved and not os.path.exists(self._compile_artifacts):
-            artifacts = torch.compiler.save_cache_artifacts()
-            if artifacts is not None:
-                with open(self._compile_artifacts, "wb") as f:
-                    f.write(artifacts[0])
-                logger.info(f"Saved compile artifacts to {self._compile_artifacts!r}")
-            self._artifacts_saved = True
-
 
 _engine = None
 _engine_lock = threading.Lock()