cache-dit 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

cache_dit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.4'
-__version_tuple__ = version_tuple = (0, 2, 4)
+__version__ = version = '0.2.6'
+__version_tuple__ = version_tuple = (0, 2, 6)

cache_dit/cache_factory/dual_block_cache/cache_context.py

@@ -71,12 +71,19 @@ class DBCacheContext:
     taylorseer_kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
     taylorseer: Optional[TaylorSeer] = None
     encoder_tarlorseer: Optional[TaylorSeer] = None
+
     # Support do_separate_classifier_free_guidance, such as Wan 2.1
     # For model that fused CFG and non-CFG into single forward step,
     # should set do_separate_classifier_free_guidance as False. For
-    # example: CogVideoX
+    # example: CogVideoX, HunyuanVideo, Mochi.
     do_separate_classifier_free_guidance: bool = False
+    # Compute cfg forward first or not, default False, namely,
+    # 0, 2, 4, ..., -> non-CFG step; 1, 3, 5, ... -> CFG step.
     cfg_compute_first: bool = False
+    # Compute spearate diff values for CFG and non-CFG step,
+    # default True. If False, we will use the computed diff from
+    # current non-CFG transformer step for current CFG step.
+    cfg_diff_compute_separate: bool = True
     cfg_taylorseer: Optional[TaylorSeer] = None
     cfg_encoder_taylorseer: Optional[TaylorSeer] = None
 
@@ -99,11 +106,17 @@ class DBCacheContext:
 
     @torch.compiler.disable
     def __post_init__(self):
+        # Some checks for settings
         if self.do_separate_classifier_free_guidance:
             assert self.enable_alter_cache is False, (
                 "enable_alter_cache must set as False if "
                 "do_separate_classifier_free_guidance is enabled."
             )
+            if self.cfg_diff_compute_separate:
+                assert self.cfg_compute_first is False, (
+                    "cfg_compute_first must set as False if "
+                    "cfg_diff_compute_separate is enabled."
+                )
 
         if "warmup_steps" not in self.taylorseer_kwargs:
             # If warmup_steps is not set in taylorseer_kwargs,
@@ -185,10 +198,17 @@ class DBCacheContext:
         # current step: incr step - 1
         self.transformer_executed_steps += 1
         if not self.do_separate_classifier_free_guidance:
-            self.executed_steps = self.transformer_executed_steps
+            self.executed_steps += 1
         else:
-            # 0,1 -> 0, 2,3 -> 1, ...
-            self.executed_steps = self.transformer_executed_steps // 2
+            # 0,1 -> 0 + 1, 2,3 -> 1 + 1, ...
+            if not self.cfg_compute_first:
+                if not self.is_separate_classifier_free_guidance_step():
+                    # transformer step: 0,2,4,...
+                    self.executed_steps += 1
+            else:
+                if self.is_separate_classifier_free_guidance_step():
+                    # transformer step: 0,2,4,...
+                    self.executed_steps += 1
 
         if not self.enable_alter_cache:
             # 0 F 1 T 2 F 3 T 4 F 5 T ...
@@ -253,6 +273,7 @@ class DBCacheContext:
 
     @torch.compiler.disable
     def add_residual_diff(self, diff):
+        # step: executed_steps - 1, not transformer_steps - 1
         step = str(self.get_current_step())
         # Only add the diff if it is not already recorded for this step
         if not self.is_separate_classifier_free_guidance_step():
@@ -299,9 +320,9 @@ class DBCacheContext:
             return False
         if self.cfg_compute_first:
             # CFG steps: 0, 2, 4, 6, ...
-            return self.get_current_transformer_step() % 2
+            return self.get_current_transformer_step() % 2 == 0
         # CFG steps: 1, 3, 5, 7, ...
-        return not self.get_current_transformer_step() % 2
+        return self.get_current_transformer_step() % 2 != 0
 
     @torch.compiler.disable
     def is_in_warmup(self):
@@ -350,6 +371,28 @@ def get_current_step():
     return cache_context.get_current_step()
 
 
+@torch.compiler.disable
+def get_current_step_residual_diff():
+    cache_context = get_current_cache_context()
+    assert cache_context is not None, "cache_context must be set before"
+    step = str(get_current_step())
+    residual_diffs = get_residual_diffs()
+    if step in residual_diffs:
+        return residual_diffs[step]
+    return None
+
+
+@torch.compiler.disable
+def get_current_step_cfg_residual_diff():
+    cache_context = get_current_cache_context()
+    assert cache_context is not None, "cache_context must be set before"
+    step = str(get_current_step())
+    cfg_residual_diffs = get_cfg_residual_diffs()
+    if step in cfg_residual_diffs:
+        return cfg_residual_diffs[step]
+    return None
+
+
 @torch.compiler.disable
 def get_current_transformer_step():
     cache_context = get_current_cache_context()
@@ -586,6 +629,13 @@ def is_separate_classifier_free_guidance_step():
     return cache_context.is_separate_classifier_free_guidance_step()
 
 
+@torch.compiler.disable
+def cfg_diff_compute_separate():
+    cache_context = get_current_cache_context()
+    assert cache_context is not None, "cache_context must be set before"
+    return cache_context.cfg_diff_compute_separate
+
+
 _current_cache_context: DBCacheContext = None
 
 
@@ -686,38 +736,49 @@ def are_two_tensors_similar(
         add_residual_diff(-2.0)
         return False
 
-    # Find the most significant token through t1 and t2, and
-    # consider the diff of the significant token. The more significant,
-    # the more important.
-    condition_thresh = get_important_condition_threshold()
-    if condition_thresh > 0.0:
-        raw_diff = (t1 - t2).abs()  # [B, seq_len, d]
-        token_m_df = raw_diff.mean(dim=-1)  # [B, seq_len]
-        token_m_t1 = t1.abs().mean(dim=-1)  # [B, seq_len]
-        # D = (t1 - t2) / t1 = 1 - (t2 / t1), if D = 0, then t1 = t2.
-        token_diff = token_m_df / token_m_t1  # [B, seq_len]
-        condition = token_diff > condition_thresh  # [B, seq_len]
-        if condition.sum() > 0:
-            condition = condition.unsqueeze(-1)  # [B, seq_len, 1]
-            condition = condition.expand_as(raw_diff)  # [B, seq_len, d]
-            mean_diff = raw_diff[condition].mean()
-            mean_t1 = t1[condition].abs().mean()
+    if all(
+        (
+            do_separate_classifier_free_guidance(),
+            is_separate_classifier_free_guidance_step(),
+            not cfg_diff_compute_separate(),
+            get_current_step_residual_diff() is not None,
+        )
+    ):
+        # Reuse computed diff value from non-CFG step
+        diff = get_current_step_residual_diff()
+    else:
+        # Find the most significant token through t1 and t2, and
+        # consider the diff of the significant token. The more significant,
+        # the more important.
+        condition_thresh = get_important_condition_threshold()
+        if condition_thresh > 0.0:
+            raw_diff = (t1 - t2).abs()  # [B, seq_len, d]
+            token_m_df = raw_diff.mean(dim=-1)  # [B, seq_len]
+            token_m_t1 = t1.abs().mean(dim=-1)  # [B, seq_len]
+            # D = (t1 - t2) / t1 = 1 - (t2 / t1), if D = 0, then t1 = t2.
+            token_diff = token_m_df / token_m_t1  # [B, seq_len]
+            condition = token_diff > condition_thresh  # [B, seq_len]
+            if condition.sum() > 0:
+                condition = condition.unsqueeze(-1)  # [B, seq_len, 1]
+                condition = condition.expand_as(raw_diff)  # [B, seq_len, d]
+                mean_diff = raw_diff[condition].mean()
+                mean_t1 = t1[condition].abs().mean()
+            else:
+                mean_diff = (t1 - t2).abs().mean()
+                mean_t1 = t1.abs().mean()
         else:
+            # Use the mean of the absolute difference of the tensors
             mean_diff = (t1 - t2).abs().mean()
             mean_t1 = t1.abs().mean()
-    else:
-        # Use the mean of the absolute difference of the tensors
-        mean_diff = (t1 - t2).abs().mean()
-        mean_t1 = t1.abs().mean()
 
-    if parallelized:
-        mean_diff = DP.all_reduce_sync(mean_diff, "avg")
-        mean_t1 = DP.all_reduce_sync(mean_t1, "avg")
+        if parallelized:
+            mean_diff = DP.all_reduce_sync(mean_diff, "avg")
+            mean_t1 = DP.all_reduce_sync(mean_t1, "avg")
 
-    # D = (t1 - t2) / t1 = 1 - (t2 / t1), if D = 0, then t1 = t2.
-    # Futher, if we assume that (H(t,  0) - H(t-1,0)) ~ 0, then,
-    # H(t-1,n) ~ H(t  ,n), which means the hidden states are similar.
-    diff = (mean_diff / mean_t1).item()
+        # D = (t1 - t2) / t1 = 1 - (t2 / t1), if D = 0, then t1 = t2.
+        # Futher, if we assume that (H(t,  0) - H(t-1,0)) ~ 0, then,
+        # H(t-1,n) ~ H(t  ,n), which means the hidden states are similar.
+        diff = (mean_diff / mean_t1).item()
 
     if logger.isEnabledFor(logging.DEBUG):
         logger.debug(f"{prefix}, diff: {diff:.6f}, threshold: {threshold:.6f}")
@@ -727,6 +788,26 @@ def are_two_tensors_similar(
     return diff < threshold
 
 
+@torch.compiler.disable
+def _debugging_set_buffer(prefix):
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug(
+            f"set {prefix}, "
+            f"transformer step: {get_current_transformer_step()}, "
+            f"executed step: {get_current_step()}"
+        )
+
+
+@torch.compiler.disable
+def _debugging_get_buffer(prefix):
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug(
+            f"get {prefix}, "
+            f"transformer step: {get_current_transformer_step()}, "
+            f"executed step: {get_current_step()}"
+        )
+
+
 # Fn buffers
 @torch.compiler.disable
 def set_Fn_buffer(buffer: torch.Tensor, prefix: str = "Fn"):
@@ -737,30 +818,38 @@ def set_Fn_buffer(buffer: torch.Tensor, prefix: str = "Fn"):
         buffer = buffer[..., ::downsample_factor]
         buffer = buffer.contiguous()
     if is_separate_classifier_free_guidance_step():
+        _debugging_set_buffer(f"{prefix}_buffer_cfg")
         set_buffer(f"{prefix}_buffer_cfg", buffer)
     else:
+        _debugging_set_buffer(f"{prefix}_buffer")
         set_buffer(f"{prefix}_buffer", buffer)
 
 
 @torch.compiler.disable
 def get_Fn_buffer(prefix: str = "Fn"):
     if is_separate_classifier_free_guidance_step():
+        _debugging_get_buffer(f"{prefix}_buffer_cfg")
         return get_buffer(f"{prefix}_buffer_cfg")
+    _debugging_get_buffer(f"{prefix}_buffer")
     return get_buffer(f"{prefix}_buffer")
 
 
 @torch.compiler.disable
 def set_Fn_encoder_buffer(buffer: torch.Tensor, prefix: str = "Fn"):
     if is_separate_classifier_free_guidance_step():
+        _debugging_set_buffer(f"{prefix}_encoder_buffer_cfg")
         set_buffer(f"{prefix}_encoder_buffer_cfg", buffer)
     else:
+        _debugging_set_buffer(f"{prefix}_encoder_buffer")
         set_buffer(f"{prefix}_encoder_buffer", buffer)
 
 
 @torch.compiler.disable
 def get_Fn_encoder_buffer(prefix: str = "Fn"):
     if is_separate_classifier_free_guidance_step():
+        _debugging_get_buffer(f"{prefix}_encoder_buffer_cfg")
         return get_buffer(f"{prefix}_encoder_buffer_cfg")
+    _debugging_get_buffer(f"{prefix}_encoder_buffer")
     return get_buffer(f"{prefix}_encoder_buffer")
 
 
@@ -786,13 +875,17 @@ def set_Bn_buffer(buffer: torch.Tensor, prefix: str = "Bn"):
                     "Falling back to default buffer retrieval."
                 )
             if is_separate_classifier_free_guidance_step():
+                _debugging_set_buffer(f"{prefix}_buffer_cfg")
                 set_buffer(f"{prefix}_buffer_cfg", buffer)
             else:
+                _debugging_set_buffer(f"{prefix}_buffer")
                 set_buffer(f"{prefix}_buffer", buffer)
     else:
         if is_separate_classifier_free_guidance_step():
+            _debugging_set_buffer(f"{prefix}_buffer_cfg")
             set_buffer(f"{prefix}_buffer_cfg", buffer)
         else:
+            _debugging_set_buffer(f"{prefix}_buffer")
             set_buffer(f"{prefix}_buffer", buffer)
 
 
@@ -815,11 +908,15 @@ def get_Bn_buffer(prefix: str = "Bn"):
                 )
             # Fallback to default buffer retrieval
             if is_separate_classifier_free_guidance_step():
+                _debugging_get_buffer(f"{prefix}_buffer_cfg")
                 return get_buffer(f"{prefix}_buffer_cfg")
+            _debugging_get_buffer(f"{prefix}_buffer")
             return get_buffer(f"{prefix}_buffer")
     else:
         if is_separate_classifier_free_guidance_step():
+            _debugging_get_buffer(f"{prefix}_buffer_cfg")
             return get_buffer(f"{prefix}_buffer_cfg")
+        _debugging_get_buffer(f"{prefix}_buffer")
         return get_buffer(f"{prefix}_buffer")
 
 
@@ -843,13 +940,17 @@ def set_Bn_encoder_buffer(buffer: torch.Tensor, prefix: str = "Bn"):
                     "Falling back to default buffer retrieval."
                 )
             if is_separate_classifier_free_guidance_step():
+                _debugging_set_buffer(f"{prefix}_encoder_buffer_cfg")
                 set_buffer(f"{prefix}_encoder_buffer_cfg", buffer)
             else:
+                _debugging_set_buffer(f"{prefix}_encoder_buffer")
                 set_buffer(f"{prefix}_encoder_buffer", buffer)
     else:
         if is_separate_classifier_free_guidance_step():
+            _debugging_set_buffer(f"{prefix}_encoder_buffer_cfg")
             set_buffer(f"{prefix}_encoder_buffer_cfg", buffer)
         else:
+            _debugging_set_buffer(f"{prefix}_encoder_buffer")
             set_buffer(f"{prefix}_encoder_buffer", buffer)
 
 
@@ -872,11 +973,15 @@ def get_Bn_encoder_buffer(prefix: str = "Bn"):
                 )
             # Fallback to default buffer retrieval
             if is_separate_classifier_free_guidance_step():
+                _debugging_get_buffer(f"{prefix}_encoder_buffer_cfg")
                 return get_buffer(f"{prefix}_encoder_buffer_cfg")
+            _debugging_get_buffer(f"{prefix}_encoder_buffer")
             return get_buffer(f"{prefix}_encoder_buffer")
     else:
         if is_separate_classifier_free_guidance_step():
+            _debugging_get_buffer(f"{prefix}_encoder_buffer_cfg")
             return get_buffer(f"{prefix}_encoder_buffer_cfg")
+        _debugging_get_buffer(f"{prefix}_encoder_buffer")
         return get_buffer(f"{prefix}_encoder_buffer")
 
 

cache_dit/cache_factory/first_block_cache/cache_context.py

@@ -370,8 +370,8 @@ def apply_prev_hidden_states_residual(
         hidden_states = hidden_states_residual + hidden_states
 
         hidden_states = hidden_states.contiguous()
-        # NOTE: We should also support taylorseer for 
-        # encoder_hidden_states approximation. Please 
+        # NOTE: We should also support taylorseer for
+        # encoder_hidden_states approximation. Please
         # use DBCache instead.
     else:
         hidden_states_residual = get_hidden_states_residual()

cache_dit/metrics/__init__.py

@@ -0,0 +1,12 @@
+from cache_dit.metrics.metrics import compute_psnr
+from cache_dit.metrics.metrics import compute_ssim
+from cache_dit.metrics.metrics import compute_mse
+from cache_dit.metrics.metrics import compute_video_psnr
+from cache_dit.metrics.metrics import compute_video_ssim
+from cache_dit.metrics.metrics import compute_video_mse
+from cache_dit.metrics.metrics import entrypoint
+from cache_dit.metrics.fid import FrechetInceptionDistance
+
+
+def main():
+    entrypoint()