[BugFix] Refresh TeaCache when num_inference_steps=None (#2240)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -265,9 +265,28 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 not getattr(req, "skip_cache_refresh", False)
                 and self.cache_backend is not None
                 and self.cache_backend.is_enabled()
-                and req.sampling_params.num_inference_steps is not None
             ):
-                self.cache_backend.refresh(self.pipeline, req.sampling_params.num_inference_steps)
+                # FIXME (Alex): When num_inference_steps is None, we defer to
+                # pipelines for default, but don't refresh the cache; the right
+                # way to do this is to merge the sampling params first.
+                #
+                # For now, if num_inference_steps is not set, we pass 0 to allow
+                # TeaCache to refresh to align with the param signature. This is
+                # okay to force refresh TeaCache because the refresh does not use
+                # num_inference_steps at all (i.e., just resets state and clears
+                # stale residuals).
+                num_inference_steps = req.sampling_params.num_inference_steps
+                if self.od_config.cache_backend == "tea_cache" and num_inference_steps is None:
+                    num_inference_steps = 0
+
+                if num_inference_steps is not None:
+                    self.cache_backend.refresh(self.pipeline, num_inference_steps)
+                else:
+                    logger.warning(
+                        "Failed to refresh the diffusion transformer cache; backend %s "
+                        "currently requires num_inference_steps to be passed explicitly",
+                        self.od_config.cache_backend,
+                    )
 
             is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
             if is_primary: