[diffusion] Refactor TeaCache

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -161,9 +161,8 @@ class SamplingParams:
 
     # TeaCache parameters
     enable_teacache: bool = False
-    teacache_params: Any = (
-        None  # TeaCacheParams or WanTeaCacheParams, set by model-specific subclass
-    )
+    cache_params: Any | None = None
+    calibrate_cache: bool = False
 
     # Profiling
     profile: bool = False
@@ -615,6 +614,12 @@ def add_argument(*name_or_flags, **kwargs):
             "--enable-teacache",
             action="store_true",
         )
+        add_argument(
+            "--calibrate-cache",
+            action="store_true",
+            default=SamplingParams.calibrate_cache,
+            help="Compute the values needed for caching.",
+        )
 
         # profiling
         add_argument(

python/sglang/multimodal_gen/test/unit/test_sampling_params.py

@@ -84,7 +84,7 @@ def coefficients_callback(_: TeaCacheParams) -> list[float]:
 
         params = SamplingParams(
             prompt="callable teacache",
-            teacache_params=TeaCacheParams(
+            cache_params=TeaCacheParams(
                 coefficients_callback=coefficients_callback,
             ),
         )