removing hardcoded magcache params values

Author: prishajain1

src/maxdiffusion/configs/base_wan_27b.yml

@@ -307,7 +307,7 @@ use_cfg_cache: False
 use_magcache: True
 magcache_thresh: 0.06
 magcache_K: 2
-retention_ratio: 0.2
+retention_ratio: 0.4
 mag_ratios_base: [1.0, 1.0, 1.00124, 1.00155, 0.99822, 0.99851, 0.99696, 0.99687, 0.99703, 0.99732, 0.9966, 0.99679, 0.99602, 0.99658, 0.99578, 0.99664, 0.99484, 0.9949, 0.99633, 0.996, 0.99659, 0.99683, 0.99534, 0.99549, 0.99584, 0.99577, 0.99681, 0.99694, 0.99563, 0.99554, 0.9944, 0.99473, 0.99594, 0.9964, 0.99466, 0.99461, 0.99453, 0.99481, 0.99389, 0.99365, 0.99391, 0.99406, 0.99354, 0.99361, 0.99283, 0.99278, 0.99268, 0.99263, 0.99057, 0.99091, 0.99125, 0.99126, 0.65523, 0.65252, 0.98808, 0.98852, 0.98765, 0.98736, 0.9851, 0.98535, 0.98311, 0.98339, 0.9805, 0.9806, 0.97776, 0.97771, 0.97278, 0.97286, 0.96731, 0.96728, 0.95857, 0.95855, 0.94385, 0.94385, 0.92118, 0.921, 0.88108, 0.88076, 0.80263, 0.80181]
 
 # SenCache: Sensitivity-Aware Caching (arXiv:2602.24208) — skip forward pass

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -296,7 +296,7 @@ guidance_scale_high: 4.0
 # The timestep threshold. If `t` is at or above this value,
 # the `high_noise_model` is considered as the required model.
 # timestep to switch between low noise and high noise transformer
-boundary_ratio: 0.875
+boundary_ratio: 0.9
 
 # Diffusion CFG cache (FasterCache-style)
 use_cfg_cache: False

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -92,10 +92,17 @@ def __call__(
       vae_only: bool = False,
       use_cfg_cache: bool = False,
       use_magcache: bool = False,
-      magcache_thresh: float = 0.12,
-      magcache_K: int = 2,
-      retention_ratio: float = 0.2,
+      magcache_thresh: Optional[float] = None,
+      magcache_K: Optional[int] = None,
+      retention_ratio: Optional[float] = None,
   ):
+    if magcache_thresh is None:
+      magcache_thresh = getattr(self.config, "magcache_thresh", 0.12)
+    if magcache_K is None:
+      magcache_K = getattr(self.config, "magcache_K", 2)
+    if retention_ratio is None:
+      retention_ratio = getattr(self.config, "retention_ratio", 0.2)
+
     if use_cfg_cache and guidance_scale <= 1.0:
       raise ValueError(
           f"use_cfg_cache=True requires guidance_scale > 1.0 (got {guidance_scale}). "

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -113,10 +113,17 @@ def __call__(
       use_cfg_cache: bool = False,
       use_sen_cache: bool = False,
       use_magcache: bool = False,
-      magcache_thresh: float = 0.04,
-      magcache_K: int = 2,
-      retention_ratio: float = 0.2,
+      magcache_thresh: Optional[float] = None,
+      magcache_K: Optional[int] = None,
+      retention_ratio: Optional[float] = None,
   ):
+    if magcache_thresh is None:
+      magcache_thresh = getattr(self.config, "magcache_thresh", 0.04)
+    if magcache_K is None:
+      magcache_K = getattr(self.config, "magcache_K", 2)
+    if retention_ratio is None:
+      retention_ratio = getattr(self.config, "retention_ratio", 0.2)
+      
     if use_cfg_cache and use_sen_cache:
       raise ValueError("use_cfg_cache and use_sen_cache are mutually exclusive. Enable only one.")
 

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -150,10 +150,17 @@ def __call__(
       output_type: Optional[str] = "np",
       rng: Optional[jax.Array] = None,
       use_magcache: bool = False,
-      magcache_thresh: float = 0.04,
-      magcache_K: int = 2,
-      retention_ratio: float = 0.2,
+      magcache_thresh: Optional[float] = None,
+      magcache_K: Optional[int] = None,
+      retention_ratio: Optional[float] = None,
   ):
+    if magcache_thresh is None:
+      magcache_thresh = getattr(self.config, "magcache_thresh", 0.04)
+    if magcache_K is None:
+      magcache_K = getattr(self.config, "magcache_K", 2)
+    if retention_ratio is None:
+      retention_ratio = getattr(self.config, "retention_ratio", 0.2)
+
     height = height or self.config.height
     width = width or self.config.width
     num_frames = num_frames or self.config.num_frames

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -168,10 +168,17 @@ def __call__(
       rng: Optional[jax.Array] = None,
       use_cfg_cache: bool = False,
       use_magcache: bool = False,
-      magcache_thresh: float = 0.04,
-      magcache_K: int = 2,
-      retention_ratio: float = 0.2,
+      magcache_thresh: Optional[float] = None,
+      magcache_K: Optional[int] = None,
+      retention_ratio: Optional[float] = None,
   ):
+    if magcache_thresh is None:
+      magcache_thresh = getattr(self.config, "magcache_thresh", 0.04)
+    if magcache_K is None:
+      magcache_K = getattr(self.config, "magcache_K", 2)
+    if retention_ratio is None:
+      retention_ratio = getattr(self.config, "retention_ratio", 0.2)
+
     if use_cfg_cache and (guidance_scale_low <= 1.0 or guidance_scale_high <= 1.0):
       raise ValueError(
           f"use_cfg_cache=True requires both guidance_scale_low > 1.0 and guidance_scale_high > 1.0 "