refactor

Author: prishajain1

src/maxdiffusion/configs/base_wan_14b.yml

@@ -332,6 +332,7 @@ use_magcache: True
 magcache_thresh: 0.12
 magcache_K: 2
 retention_ratio: 0.2
+mag_ratios_base: [1.0, 1.0, 1.02504, 1.03017, 1.00025, 1.00251, 0.9985, 0.99962, 0.99779, 0.99771, 0.9966, 0.99658, 0.99482, 0.99476, 0.99467, 0.99451, 0.99664, 0.99656, 0.99434, 0.99431, 0.99533, 0.99545, 0.99468, 0.99465, 0.99438, 0.99434, 0.99516, 0.99517, 0.99384, 0.9938, 0.99404, 0.99401, 0.99517, 0.99516, 0.99409, 0.99408, 0.99428, 0.99426, 0.99347, 0.99343, 0.99418, 0.99416, 0.99271, 0.99279, 0.99313, 0.99311, 0.99215, 0.99215, 0.99218, 0.99215, 0.99216, 0.99217, 0.99163, 0.99161, 0.99138, 0.99135, 0.98982, 0.9898, 0.98996, 0.98995, 0.9887, 0.98866, 0.98772, 0.9877, 0.98767, 0.98765, 0.98573, 0.9857, 0.98501, 0.98498, 0.9838, 0.98376, 0.98177, 0.98173, 0.98037, 0.98035, 0.97678, 0.97677, 0.97546, 0.97543, 0.97184, 0.97183, 0.96711, 0.96708, 0.96349, 0.96345, 0.95629, 0.95625, 0.94926, 0.94929, 0.93964, 0.93961, 0.92511, 0.92504, 0.90693, 0.90678, 0.8796, 0.87945, 0.86111, 0.86189]
 
 # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
 guidance_rescale: 0.0

src/maxdiffusion/configs/base_wan_27b.yml

@@ -308,6 +308,7 @@ use_magcache: True
 magcache_thresh: 0.12
 magcache_K: 2
 retention_ratio: 0.2
+mag_ratios_base: [1.0, 1.0, 1.00124, 1.00155, 0.99822, 0.99851, 0.99696, 0.99687, 0.9977, 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
 # when predicted output change (based on accumulated latent/timestep drift) is small

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -292,6 +292,8 @@ use_magcache: True
 magcache_thresh: 0.12
 magcache_K: 2
 retention_ratio: 0.2
+mag_ratios_base_720p: [1.0, 1.0, 0.99428, 0.99498, 0.98588, 0.98621, 0.98273, 0.98281, 0.99018, 0.99023, 0.98911, 0.98917, 0.98646, 0.98652, 0.99454, 0.99456, 0.9891, 0.98909, 0.99124, 0.99127, 0.99102, 0.99103, 0.99215, 0.99212, 0.99515, 0.99515, 0.99576, 0.99572, 0.99068, 0.99072, 0.99097, 0.99097, 0.99166, 0.99169, 0.99041, 0.99042, 0.99201, 0.99198, 0.99101, 0.99101, 0.98599, 0.98603, 0.98845, 0.98844, 0.98848, 0.98851, 0.98862, 0.98857, 0.98718, 0.98719, 0.98497, 0.98497, 0.98264, 0.98263, 0.98389, 0.98393, 0.97938, 0.9794, 0.97535, 0.97536, 0.97498, 0.97499, 0.973, 0.97301, 0.96827, 0.96828, 0.96261, 0.96263, 0.95335, 0.9534, 0.94649, 0.94655, 0.93397, 0.93414, 0.91636, 0.9165, 0.89088, 0.89109, 0.8679, 0.86768]
+mag_ratios_base_480p: [1.0, 1.0, 0.98783, 0.98993, 0.97559, 0.97593, 0.98311, 0.98319, 0.98202, 0.98225, 0.9888, 0.98878, 0.98762, 0.98759, 0.98957, 0.98971, 0.99052, 0.99043, 0.99383, 0.99384, 0.98857, 0.9886, 0.99065, 0.99068, 0.98845, 0.98847, 0.99057, 0.99057, 0.98957, 0.98961, 0.98601, 0.9861, 0.98823, 0.98823, 0.98756, 0.98759, 0.98808, 0.98814, 0.98721, 0.98724, 0.98571, 0.98572, 0.98543, 0.98544, 0.98157, 0.98165, 0.98411, 0.98413, 0.97952, 0.97953, 0.98149, 0.9815, 0.9774, 0.97742, 0.97825, 0.97826, 0.97355, 0.97361, 0.97085, 0.97087, 0.97056, 0.97055, 0.96588, 0.96587, 0.96113, 0.96124, 0.9567, 0.95681, 0.94961, 0.94969, 0.93973, 0.93988, 0.93217, 0.93224, 0.91878, 0.91896, 0.90955, 0.90954, 0.92617, 0.92616]
 
 # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
 guidance_rescale: 0.0

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -304,6 +304,7 @@ use_magcache: True
 magcache_thresh: 0.12
 magcache_K: 2
 retention_ratio: 0.2
+mag_ratios_base: [1.0, 1.0, 0.99512, 0.99559, 0.99559, 0.99561, 0.99595, 0.99577, 0.99512, 0.99512, 0.99546, 0.99534, 0.99543, 0.99531, 0.99496, 0.99491, 0.99504, 0.99499, 0.99444, 0.99449, 0.99481, 0.99481, 0.99435, 0.99435, 0.9943, 0.99431, 0.99411, 0.99406, 0.99373, 0.99376, 0.99413, 0.99405, 0.99363, 0.99359, 0.99335, 0.99331, 0.99244, 0.99243, 0.99229, 0.99229, 0.99239, 0.99236, 0.99163, 0.9916, 0.99149, 0.99151, 0.99191, 0.99192, 0.9898, 0.98981, 0.9899, 0.98987, 0.98849, 0.98849, 0.98846, 0.98846, 0.98861, 0.98861, 0.9874, 0.98738, 0.98588, 0.98589, 0.98539, 0.98534, 0.98444, 0.98439, 0.9831, 0.98309, 0.98119, 0.98118, 0.98001, 0.98, 0.97862, 0.97859, 0.97555, 0.97558, 0.97392, 0.97388, 0.97152, 0.97145, 0.96871, 0.9687, 0.96435, 0.96434, 0.96129, 0.96127, 0.95639, 0.95638, 0.95176, 0.95175, 0.94446, 0.94452, 0.93972, 0.93974, 0.93575, 0.9359, 0.93537, 0.93552, 0.96655, 0.96616]
 
 # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
 guidance_rescale: 0.0

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -910,4 +910,116 @@ def nearest_interp(src, target_len):
         return np.array([src[-1]])
     import numpy as np
     indices = np.round(np.linspace(0, src_len - 1, target_len)).astype(np.int32)
-    return src[indices]
+    return src[indices]
+
+def init_magcache(num_inference_steps, retention_ratio, mag_ratios_base):
+    """Initialize MagCache variables and interpolate ratios.
+    
+    Args:
+        num_inference_steps: Number of inference steps.
+        retention_ratio: Retention ratio of unchanged steps.
+        mag_ratios_base: Base magnitude ratios array or list.
+    """
+    import numpy as np
+    
+    accumulated_ratio_cond = 1.0
+    accumulated_ratio_uncond = 1.0
+    accumulated_err_cond = 0.0
+    accumulated_err_uncond = 0.0
+    accumulated_steps_cond = 0
+    accumulated_steps_uncond = 0
+    cached_residual = None
+
+    skip_warmup = int(num_inference_steps * retention_ratio)
+
+    mag_ratios_base = np.array(mag_ratios_base)
+
+    if len(mag_ratios_base) != num_inference_steps * 2:
+        mag_cond = nearest_interp(mag_ratios_base[0::2], num_inference_steps)
+        mag_uncond = nearest_interp(mag_ratios_base[1::2], num_inference_steps)
+        mag_ratios = np.concatenate([mag_cond.reshape(-1, 1), mag_uncond.reshape(-1, 1)], axis=1).reshape(-1)
+    else:
+        mag_ratios = mag_ratios_base
+
+    return (
+        accumulated_ratio_cond,
+        accumulated_ratio_uncond,
+        accumulated_err_cond,
+        accumulated_err_uncond,
+        accumulated_steps_cond,
+        accumulated_steps_uncond,
+        cached_residual,
+        skip_warmup,
+        mag_ratios,
+    )
+
+def magcache_step(
+    step,
+    mag_ratios,
+    accumulated_state,
+    magcache_thresh,
+    magcache_K,
+    skip_warmup,
+):
+    """Update MagCache accumulated state and decide if to skip.
+    
+    Args:
+        step: Current inference step.
+        mag_ratios: Interpolated magnitude ratios array.
+        accumulated_state: Tuple containing accumulated variables.
+        magcache_thresh: Error threshold.
+        magcache_K: Max skip steps.
+        skip_warmup: Warmup steps threshold.
+    """
+    import numpy as np
+    
+    (
+        accumulated_ratio_cond,
+        accumulated_ratio_uncond,
+        accumulated_err_cond,
+        accumulated_err_uncond,
+        accumulated_steps_cond,
+        accumulated_steps_uncond,
+    ) = accumulated_state
+
+    cur_mag_ratio_cond = mag_ratios[step * 2]
+    cur_mag_ratio_uncond = mag_ratios[step * 2 + 1]
+
+    skip_blocks = False
+    if step >= skip_warmup:
+        new_ratio_cond = accumulated_ratio_cond * cur_mag_ratio_cond
+        new_ratio_uncond = accumulated_ratio_uncond * cur_mag_ratio_uncond
+
+        err_cond = np.abs(1.0 - new_ratio_cond)
+        err_uncond = np.abs(1.0 - new_ratio_uncond)
+
+        if (
+            accumulated_err_cond + err_cond < magcache_thresh
+            and accumulated_steps_cond < magcache_K
+            and accumulated_err_uncond + err_uncond < magcache_thresh
+            and accumulated_steps_uncond < magcache_K
+        ):
+            skip_blocks = True
+            accumulated_ratio_cond = new_ratio_cond
+            accumulated_ratio_uncond = new_ratio_uncond
+            accumulated_err_cond += err_cond
+            accumulated_err_uncond += err_uncond
+            accumulated_steps_cond += 1
+            accumulated_steps_uncond += 1
+        else:
+            accumulated_ratio_cond = 1.0
+            accumulated_ratio_uncond = 1.0
+            accumulated_err_cond = 0.0
+            accumulated_err_uncond = 0.0
+            accumulated_steps_cond = 0
+            accumulated_steps_uncond = 0
+
+    new_state = (
+        accumulated_ratio_cond,
+        accumulated_ratio_uncond,
+        accumulated_err_cond,
+        accumulated_err_uncond,
+        accumulated_steps_cond,
+        accumulated_steps_uncond,
+    )
+    return skip_blocks, new_state

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .wan_pipeline import WanPipeline, transformer_forward_pass, transformer_forward_pass_full_cfg, transformer_forward_pass_cfg_cache, nearest_interp
+from .wan_pipeline import WanPipeline, transformer_forward_pass, transformer_forward_pass_full_cfg, transformer_forward_pass_cfg_cache, nearest_interp, init_magcache, magcache_step
 from ...models.wan.transformers.transformer_wan import WanModel
 from typing import List, Union, Optional, Any
 from ...pyconfig import HyperParameters
@@ -131,6 +131,7 @@ def __call__(
         magcache_K=magcache_K,
         retention_ratio=retention_ratio,
         height=height,
+        mag_ratios_base=self.config.mag_ratios_base if hasattr(self.config, "mag_ratios_base") else None,
     )
 
     with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
@@ -163,6 +164,7 @@ def run_inference_2_1(
     magcache_K: int = 2,
     retention_ratio: float = 0.2,
     height: int = 480,
+    mag_ratios_base: Optional[List[float]] = None,
 ):
   """Denoising loop for WAN 2.1 T2V with FasterCache CFG-Cache.
 
@@ -233,59 +235,42 @@ def run_inference_2_1(
   cached_noise_cond = None
   cached_noise_uncond = None
 
-  if use_magcache and do_classifier_free_guidance:
-    # ── MagCache Execution Path ──
-    accumulated_ratio_cond = 1.0
-    accumulated_ratio_uncond = 1.0
-    accumulated_err_cond = 0.0
-    accumulated_err_uncond = 0.0
-    accumulated_steps_cond = 0
-    accumulated_steps_uncond = 0
-    cached_residual = None
-
-    skip_warmup = int(num_inference_steps * retention_ratio)
-
-    # 14B Ratios
-    mag_ratios_base = np.array([1.0]*2+[1.02504, 1.03017, 1.00025, 1.00251, 0.9985, 0.99962, 0.99779, 0.99771, 0.9966, 0.99658, 0.99482, 0.99476, 0.99467, 0.99451, 0.99664, 0.99656, 0.99434, 0.99431, 0.99533, 0.99545, 0.99468, 0.99465, 0.99438, 0.99434, 0.99516, 0.99517, 0.99384, 0.9938, 0.99404, 0.99401, 0.99517, 0.99516, 0.99409, 0.99408, 0.99428, 0.99426, 0.99347, 0.99343, 0.99418, 0.99416, 0.99271, 0.99269, 0.99313, 0.99311, 0.99215, 0.99215, 0.99218, 0.99215, 0.99216, 0.99217, 0.99163, 0.99161, 0.99138, 0.99135, 0.98982, 0.9898, 0.98996, 0.98995, 0.9887, 0.98866, 0.98772, 0.9877, 0.98767, 0.98765, 0.98573, 0.9857, 0.98501, 0.98498, 0.9838, 0.98376, 0.98177, 0.98173, 0.98037, 0.98035, 0.97678, 0.97677, 0.97546, 0.97543, 0.97184, 0.97183, 0.96711, 0.96708, 0.96349, 0.96345, 0.95629, 0.95625, 0.94926, 0.94929, 0.93964, 0.93961, 0.92511, 0.92504, 0.90693, 0.90678, 0.8796, 0.87945, 0.86111, 0.86189])
-
-    if len(mag_ratios_base) != num_inference_steps * 2:
-      mag_cond = nearest_interp(mag_ratios_base[0::2], num_inference_steps)
-      mag_uncond = nearest_interp(mag_ratios_base[1::2], num_inference_steps)
-      mag_ratios = np.concatenate([mag_cond.reshape(-1, 1), mag_uncond.reshape(-1, 1)], axis=1).reshape(-1)
-    else:
-      mag_ratios = mag_ratios_base
+  if use_magcache and do_cfg:
+    (
+        accumulated_ratio_cond,
+        accumulated_ratio_uncond,
+        accumulated_err_cond,
+        accumulated_err_uncond,
+        accumulated_steps_cond,
+        accumulated_steps_uncond,
+        cached_residual,
+        skip_warmup,
+        mag_ratios,
+    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
 
     for step in range(num_inference_steps):
       t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
       timestep = jnp.broadcast_to(t, bsz * 2 if do_cfg else bsz)
 
-      skip_blocks = False
-      if step >= skip_warmup:
-        ratio_cond = mag_ratios[step * 2]
-        ratio_uncond = mag_ratios[step * 2 + 1]
-
-        new_ratio_cond = accumulated_ratio_cond * ratio_cond
-        new_ratio_uncond = accumulated_ratio_uncond * ratio_uncond
-
-        err_cond = np.abs(1.0 - new_ratio_cond)
-        err_uncond = np.abs(1.0 - new_ratio_uncond)
-
-        if (accumulated_err_cond + err_cond < magcache_thresh and accumulated_steps_cond < magcache_K and
-            accumulated_err_uncond + err_uncond < magcache_thresh and accumulated_steps_uncond < magcache_K):
-          skip_blocks = True
-          accumulated_ratio_cond = new_ratio_cond
-          accumulated_ratio_uncond = new_ratio_uncond
-          accumulated_err_cond += err_cond
-          accumulated_err_uncond += err_uncond
-          accumulated_steps_cond += 1
-          accumulated_steps_uncond += 1
-        else:
-          accumulated_ratio_cond = 1.0
-          accumulated_ratio_uncond = 1.0
-          accumulated_err_cond = 0.0
-          accumulated_err_uncond = 0.0
-          accumulated_steps_cond = 0
-          accumulated_steps_uncond = 0
+      accumulated_state = (
+          accumulated_ratio_cond,
+          accumulated_ratio_uncond,
+          accumulated_err_cond,
+          accumulated_err_uncond,
+          accumulated_steps_cond,
+          accumulated_steps_uncond,
+      )
+      skip_blocks, accumulated_state = magcache_step(
+          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup
+      )
+      (
+          accumulated_ratio_cond,
+          accumulated_ratio_uncond,
+          accumulated_err_cond,
+          accumulated_err_uncond,
+          accumulated_steps_cond,
+          accumulated_steps_uncond,
+      ) = accumulated_state
 
       outputs = transformer_forward_pass(
           graphdef,

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .wan_pipeline import WanPipeline, transformer_forward_pass, transformer_forward_pass_full_cfg, transformer_forward_pass_cfg_cache, nearest_interp
+from .wan_pipeline import WanPipeline, transformer_forward_pass, transformer_forward_pass_full_cfg, transformer_forward_pass_cfg_cache, nearest_interp, init_magcache, magcache_step
 from ...models.wan.transformers.transformer_wan import WanModel
 from typing import List, Union, Optional
 from ...pyconfig import HyperParameters
@@ -169,6 +169,7 @@ def __call__(
         magcache_K=magcache_K,
         retention_ratio=retention_ratio,
         height=height,
+        mag_ratios_base=self.config.mag_ratios_base if hasattr(self.config, "mag_ratios_base") else None,
     )
 
     with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
@@ -210,6 +211,7 @@ def run_inference_2_2(
     magcache_K: int = 2,
     retention_ratio: float = 0.2,
     height: int = 480,
+    mag_ratios_base: Optional[List[float]] = None,
 ):
   """Denoising loop for WAN 2.2 T2V with optional caching acceleration.
 
@@ -452,59 +454,44 @@ def run_inference_2_2(
 
   # ── MagCache path ──
   if use_magcache and do_classifier_free_guidance:
-    accumulated_ratio_cond = 1.0
-    accumulated_ratio_uncond = 1.0
-    accumulated_err_cond = 0.0
-    accumulated_err_uncond = 0.0
-    accumulated_steps_cond = 0
-    accumulated_steps_uncond = 0
-    cached_residual = None
-
-    skip_warmup = int(num_inference_steps * retention_ratio)
-
-    # Pre-calculated 2.2 T2V ratios
-    mag_ratios_base = np.array([1.0]*2+[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])
-
-    if len(mag_ratios_base) != num_inference_steps * 2:
-      mag_cond = nearest_interp(mag_ratios_base[0::2], num_inference_steps)
-      mag_uncond = nearest_interp(mag_ratios_base[1::2], num_inference_steps)
-      mag_ratios = np.concatenate([mag_cond.reshape(-1, 1), mag_uncond.reshape(-1, 1)], axis=1).reshape(-1)
-    else:
-      mag_ratios = mag_ratios_base
+    (
+        accumulated_ratio_cond,
+        accumulated_ratio_uncond,
+        accumulated_err_cond,
+        accumulated_err_uncond,
+        accumulated_steps_cond,
+        accumulated_steps_uncond,
+        cached_residual,
+        skip_warmup,
+        mag_ratios,
+    ) = init_magcache(num_inference_steps, retention_ratio, mag_ratios_base)
 
     timesteps_np = np.array(scheduler_state.timesteps, dtype=np.int32)
     step_uses_high = [bool(timesteps_np[s] >= boundary) for s in range(num_inference_steps)]
     prompt_embeds_combined = jnp.concatenate([prompt_embeds, negative_prompt_embeds], axis=0)
 
     for step in range(num_inference_steps):
       t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
-      cur_mag_ratio_cond = mag_ratios[step*2]
-      cur_mag_ratio_uncond = mag_ratios[step*2+1]
-
-      skip_blocks = False
-      if step >= skip_warmup:
-        new_ratio_cond = accumulated_ratio_cond * cur_mag_ratio_cond
-        new_ratio_uncond = accumulated_ratio_uncond * cur_mag_ratio_uncond
-
-        err_cond = np.abs(1.0 - new_ratio_cond)
-        err_uncond = np.abs(1.0 - new_ratio_uncond)
-
-        if (accumulated_err_cond + err_cond < magcache_thresh and accumulated_steps_cond < magcache_K and
-            accumulated_err_uncond + err_uncond < magcache_thresh and accumulated_steps_uncond < magcache_K):
-          skip_blocks = True
-          accumulated_ratio_cond = new_ratio_cond
-          accumulated_ratio_uncond = new_ratio_uncond
-          accumulated_err_cond += err_cond
-          accumulated_err_uncond += err_uncond
-          accumulated_steps_cond += 1
-          accumulated_steps_uncond += 1
-        else:
-          accumulated_ratio_cond = 1.0
-          accumulated_ratio_uncond = 1.0
-          accumulated_err_cond = 0.0
-          accumulated_err_uncond = 0.0
-          accumulated_steps_cond = 0
-          accumulated_steps_uncond = 0
+
+      accumulated_state = (
+          accumulated_ratio_cond,
+          accumulated_ratio_uncond,
+          accumulated_err_cond,
+          accumulated_err_uncond,
+          accumulated_steps_cond,
+          accumulated_steps_uncond,
+      )
+      skip_blocks, accumulated_state = magcache_step(
+          step, mag_ratios, accumulated_state, magcache_thresh, magcache_K, skip_warmup
+      )
+      (
+          accumulated_ratio_cond,
+          accumulated_ratio_uncond,
+          accumulated_err_cond,
+          accumulated_err_uncond,
+          accumulated_steps_cond,
+          accumulated_steps_uncond,
+      ) = accumulated_state
 
       if step_uses_high[step]:
         graphdef, state, rest = high_noise_graphdef, high_noise_state, high_noise_rest