formatting

Author: AlanPonnachan

src/diffusers/hooks/mag_cache.py

@@ -162,7 +162,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
 
         current_step = state.step_index
         if current_step >= len(self.config.mag_ratios):
-            # Safety fallback if steps exceed config
             current_scale = 1.0
         else:
             current_scale = self.config.mag_ratios[current_step]
@@ -174,16 +173,13 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
             state.accumulated_steps += 1
             state.accumulated_err += abs(1.0 - state.accumulated_ratio)
 
-            # Check skip condition
-            # We must have a previous residual to skip
             if (
                 state.previous_residual is not None
                 and state.accumulated_err <= self.config.threshold
                 and state.accumulated_steps <= self.config.max_skip_steps
             ):
                 should_compute = False
             else:
-                # Reset accumulators if we decide to compute (and we are past retention)
                 state.accumulated_ratio = 1.0
                 state.accumulated_steps = 0
                 state.accumulated_err = 0.0
@@ -207,7 +203,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
             ):
                 diff = output.shape[1] - res.shape[1]
                 if diff > 0:
-                    # Add residual to the end
                     output = output.clone()
                     output[:, diff:, :] = output[:, diff:, :] + res
                 else:
@@ -239,7 +234,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
                 return output
 
         else:
-            # Run original forward
             output = self.fn_ref.original_forward(*args, **kwargs)
             return output
 
@@ -302,15 +296,13 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
 
             in_hidden = state.head_block_input
 
-            # Calculate residual
             if out_hidden.shape == in_hidden.shape:
                 residual = out_hidden - in_hidden
             elif out_hidden.ndim == 3 and in_hidden.ndim == 3 and out_hidden.shape[2] == in_hidden.shape[2]:
                 diff = in_hidden.shape[1] - out_hidden.shape[1]
                 if diff == 0:
                     residual = out_hidden - in_hidden
                 else:
-                    # Fallback: Just calculate residual on matching tail (Image part usually at end)
                     residual = out_hidden - in_hidden
 
             state.previous_residual = residual