Update src/caching.py

src/caching.py

@@ -63,7 +63,8 @@ def are_two_tensors_similar(t1, t2, *, threshold=0.85):
     mean_t1 = t1.abs().mean()
     diff = mean_diff / mean_t1
     return diff.item() < threshold
-    
+
+
 class CachedTransformerBlocks(torch.nn.Module):
     def __init__(
         self,
@@ -72,7 +73,7 @@ class CachedTransformerBlocks(torch.nn.Module):
         *,
         transformer=None,
         residual_diff_threshold=0.05,
-        return_hidden_states_first=False,  # Changed default to False for Flux
+        return_hidden_states_first=False,
     ):
         super().__init__()
         self.transformer = transformer
@@ -82,82 +83,94 @@ class CachedTransformerBlocks(torch.nn.Module):
         self.return_hidden_states_first = return_hidden_states_first
 
     def forward(self, hidden_states, encoder_hidden_states, *args, **kwargs):
-        # For Flux architecture, we need to handle the order differently
-        if not self.return_hidden_states_first:
-            hidden_states, encoder_hidden_states = encoder_hidden_states, hidden_states
-
-        if self.residual_diff_threshold <= 0.0:
-            for block in self.transformer_blocks:
-                if self.return_hidden_states_first:
-                    hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, *args, **kwargs)
-                else:
-                    encoder_hidden_states, hidden_states = block(encoder_hidden_states, hidden_states, *args, **kwargs)
+        # Store original states before any transformations
+        original_hidden_states = hidden_states
+        original_encoder_hidden_states = encoder_hidden_states
 
-            return (hidden_states, encoder_hidden_states) if self.return_hidden_states_first else (encoder_hidden_states, hidden_states)
-
-        original_encoder_states = encoder_hidden_states
+        # Process first block
         first_block = self.transformer_blocks[0]
+        hidden_states, encoder_hidden_states = first_block(
+            hidden_states, encoder_hidden_states, *args, **kwargs
+        )
 
-        if self.return_hidden_states_first:
-            hidden_states, encoder_hidden_states = first_block(hidden_states, encoder_hidden_states, *args, **kwargs)
-            first_residual = hidden_states - original_encoder_states
-        else:
-            encoder_hidden_states, hidden_states = first_block(encoder_hidden_states, hidden_states, *args, **kwargs)
-            first_residual = encoder_hidden_states - original_encoder_states
+        # Calculate residual from first block
+        first_hidden_states_residual = hidden_states - original_hidden_states
 
         cache_context = get_current_cache_context()
-        prev_residual = cache_context.get_buffer("first_residual")
+        prev_residual = cache_context.get_buffer("first_hidden_states_residual")
+
         can_use_cache = prev_residual is not None and are_two_tensors_similar(
-            prev_residual, first_residual, threshold=self.residual_diff_threshold
+            prev_residual,
+            first_hidden_states_residual,
+            threshold=self.residual_diff_threshold
         )
 
         if can_use_cache:
-            residual = cache_context.get_buffer("residual")
-            if self.return_hidden_states_first:
-                hidden_states = residual + hidden_states
-            else:
-                encoder_hidden_states = residual + encoder_hidden_states
+            # Use cached residuals
+            hidden_states_residual = cache_context.get_buffer("hidden_states_residual")
+            encoder_hidden_states_residual = cache_context.get_buffer("encoder_hidden_states_residual")
+
+            hidden_states = hidden_states + hidden_states_residual
+            encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_residual
         else:
-            cache_context.set_buffer("first_residual", first_residual)
-            original_states = original_encoder_states
+            # Process remaining blocks and cache results
+            cache_context.set_buffer("first_hidden_states_residual", first_hidden_states_residual)
 
             for block in self.transformer_blocks[1:]:
-                if self.return_hidden_states_first:
-                    hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, *args, **kwargs)
-                else:
-                    encoder_hidden_states, hidden_states = block(encoder_hidden_states, hidden_states, *args, **kwargs)
-
-            if self.return_hidden_states_first:
-                cache_context.set_buffer("residual", hidden_states - original_states)
-            else:
-                cache_context.set_buffer("residual", encoder_hidden_states - original_states)
-
-        return (hidden_states, encoder_hidden_states) if self.return_hidden_states_first else (encoder_hidden_states, hidden_states)
-
-def apply_cache_on_transformer(transformer: FluxTransformer2DModel, *, residual_diff_threshold=0.05):
-    cached_blocks = torch.nn.ModuleList([
+                hidden_states, encoder_hidden_states = block(
+                    hidden_states, encoder_hidden_states, *args, **kwargs
+                )
+
+            if self.single_transformer_blocks is not None:
+                for block in self.single_transformer_blocks:
+                    hidden_states = block(hidden_states, *args, **kwargs)
+
+            # Store residuals for future use
+            cache_context.set_buffer(
+                "hidden_states_residual",
+                hidden_states - original_hidden_states
+            )
+            cache_context.set_buffer(
+                "encoder_hidden_states_residual",
+                encoder_hidden_states - original_encoder_hidden_states
+            )
+
+        return hidden_states, encoder_hidden_states
+
+def apply_cache_on_transformer(
+    transformer: FluxTransformer2DModel,
+    *,
+    residual_diff_threshold=0.05,
+):
+    cached_transformer_blocks = torch.nn.ModuleList([
         CachedTransformerBlocks(
             transformer.transformer_blocks,
             transformer.single_transformer_blocks if hasattr(transformer, 'single_transformer_blocks') else None,
             transformer=transformer,
             residual_diff_threshold=residual_diff_threshold,
-            return_hidden_states_first=False  # Specifically for Flux
         )
     ])
+    dummy_single_transformer_blocks = torch.nn.ModuleList()
 
     original_forward = transformer.forward
 
-    @functools.wraps(transformer.__class__.forward)
+    @functools.wraps(original_forward)
     def new_forward(self, *args, **kwargs):
-        with unittest.mock.patch.object(self, "transformer_blocks", cached_blocks):
-            if hasattr(self, 'single_transformer_blocks'):
-                with unittest.mock.patch.object(self, "single_transformer_blocks", torch.nn.ModuleList()):
-                    return original_forward(*args, **kwargs)
+        with unittest.mock.patch.object(
+            self,
+            "transformer_blocks",
+            cached_transformer_blocks,
+        ), unittest.mock.patch.object(
+            self,
+            "single_transformer_blocks",
+            dummy_single_transformer_blocks,
+        ):
             return original_forward(*args, **kwargs)
 
     transformer.forward = new_forward.__get__(transformer)
     return transformer
     
+    
 def apply_cache_on_pipe(pipe: DiffusionPipeline, *, shallow_patch: bool = False, **kwargs):
     original_call = pipe.__class__.__call__