update

.ai/AGENTS.md

@@ -1,77 +0,0 @@
-# Diffusers — Agent Guide
-
-## Coding style
-
-Strive to write code as simple and explicit as possible.
-
-- Minimize small helper/utility functions — inline the logic instead. A reader should be able to follow the full flow without jumping between functions.
-- No defensive code or unused code paths — do not add fallback paths, safety checks, or configuration options "just in case". When porting from a research repo, delete training-time code paths, experimental flags, and ablation branches entirely — only keep the inference path you are actually integrating.
-- Do not guess user intent and silently correct behavior. Make the expected inputs clear in the docstring, and raise a concise error for unsupported cases rather than adding complex fallback logic.
-
----
-
-### Dependencies
-- No new mandatory dependency without discussion (e.g. `einops`)
-- Optional deps guarded with `is_X_available()` and a dummy in `utils/dummy_*.py`
-
-## Code formatting
-- `make style` and `make fix-copies` should be run as the final step before opening a PR
-
-### Copied Code
-- Many classes are kept in sync with a source via a `# Copied from ...` header comment
-- Do not edit a `# Copied from` block directly — run `make fix-copies` to propagate changes from the source
-- Remove the header to intentionally break the link
-
-### Models
-- All layer calls should be visible directly in `forward` — avoid helper functions that hide `nn.Module` calls.
-- Try to not introduce graph breaks as much as possible for better compatibility with `torch.compile`. For example, DO NOT arbitrarily insert operations from NumPy in the forward implementations.
-- Attention must follow the diffusers pattern: both the `Attention` class and its processor are defined in the model file. The processor's `__call__` handles the actual compute and must use `dispatch_attention_fn` rather than calling `F.scaled_dot_product_attention` directly. The attention class inherits `AttentionModuleMixin` and declares `_default_processor_cls` and `_available_processors`.
-
-```python
-# transformer_mymodel.py
-
-class MyModelAttnProcessor:
-    _attention_backend = None
-    _parallel_config = None
-
-    def __call__(self, attn, hidden_states, attention_mask=None, ...):
-        query = attn.to_q(hidden_states)
-        key = attn.to_k(hidden_states)
-        value = attn.to_v(hidden_states)
-        # reshape, apply rope, etc.
-        hidden_states = dispatch_attention_fn(
-            query, key, value,
-            attn_mask=attention_mask,
-            backend=self._attention_backend,
-            parallel_config=self._parallel_config,
-        )
-        hidden_states = hidden_states.flatten(2, 3)
-        return attn.to_out[0](hidden_states)
-
-
-class MyModelAttention(nn.Module, AttentionModuleMixin):
-    _default_processor_cls = MyModelAttnProcessor
-    _available_processors = [MyModelAttnProcessor]
-
-    def __init__(self, query_dim, heads=8, dim_head=64, ...):
-        super().__init__()
-        self.to_q = nn.Linear(query_dim, heads * dim_head, bias=False)
-        self.to_k = nn.Linear(query_dim, heads * dim_head, bias=False)
-        self.to_v = nn.Linear(query_dim, heads * dim_head, bias=False)
-        self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, query_dim), nn.Dropout(0.0)])
-        self.set_processor(MyModelAttnProcessor())
-
-    def forward(self, hidden_states, attention_mask=None, **kwargs):
-        return self.processor(self, hidden_states, attention_mask, **kwargs)
-```
-
-Consult the implementations in `src/diffusers/models/transformers/` if you need further references.
-
-### Pipeline
-- All pipelines must inherit from `DiffusionPipeline`. Consult implementations in `src/diffusers/pipelines` in case you need references.
-- DO NOT use an existing pipeline class (e.g., `FluxPipeline`) to override another pipeline (e.g., `FluxImg2ImgPipeline` which will be a part of the core codebase (`src`).
-
-
-### Tests
-- Slow tests gated with `@slow` and `RUN_SLOW=1`
-- All model-level tests must use the `BaseModelTesterConfig`, `ModelTesterMixin`, `MemoryTesterMixin`, `AttentionTesterMixin`, `LoraTesterMixin`, and `TrainingTesterMixin` classes initially to write the tests. Any additional tests should be added after discussions with the maintainers. Use `tests/models/transformers/test_models_transformer_flux.py` as a reference.

.gitignore

@@ -178,8 +178,4 @@ tags
 .ruff_cache
 
 # wandb
-wandb
-
-# AI agent generated symlinks
-/AGENTS.md
-/CLAUDE.md
+wandb

Makefile

@@ -1,4 +1,4 @@
-.PHONY: deps_table_update modified_only_fixup extra_style_checks quality style fixup fix-copies test test-examples codex claude clean-ai
+.PHONY: deps_table_update modified_only_fixup extra_style_checks quality style fixup fix-copies test test-examples
 
 # make sure to test the local checkout in scripts and not the pre-installed one (don't use quotes!)
 export PYTHONPATH = src
@@ -98,14 +98,3 @@ post-release:
 
 post-patch:
 	python utils/release.py --post_release --patch
-
-# AI agent symlinks
-
-codex:
-	ln -snf .ai/AGENTS.md AGENTS.md
-
-claude:
-	ln -snf .ai/AGENTS.md CLAUDE.md
-
-clean-ai:
-	rm -f AGENTS.md CLAUDE.md

docs/source/en/conceptual/contribution.md

@@ -565,16 +565,4 @@ $ git push --set-upstream origin your-branch-for-syncing
 
 ### Style guide
 
-For documentation strings, 🧨 Diffusers follows the [Google style](https://google.github.io/styleguide/pyguide.html).
-
-
-## Coding with AI agents
-
-The repository keeps AI-agent configuration in `.ai/` and exposes local agent files via symlinks.
-
-- **Source of truth** — edit `.ai/AGENTS.md` (and any future `.ai/skills/`)
-- **Don't edit** generated root-level `AGENTS.md` or `CLAUDE.md` — they are symlinks
-- Setup commands:
-  - `make codex` — symlink for OpenAI Codex
-  - `make claude` — symlink for Claude Code
-  - `make clean-ai` — remove generated symlinks
+For documentation strings, 🧨 Diffusers follows the [Google style](https://google.github.io/styleguide/pyguide.html).

src/diffusers/modular_pipelines/modular_pipeline_utils.py

@@ -309,16 +309,16 @@ class ComponentSpec:
                 f"`type_hint` is required when loading a single file model but is missing for component: {self.name}"
             )
 
-        from diffusers import AutoModel
-
         # `torch_dtype` is not an accepted parameter for tokenizers and processors.
         # As a result, it gets stored in `init_kwargs`, which are written to the config
         # during save. This causes JSON serialization to fail when saving the component.
-        if self.type_hint is not None and not issubclass(self.type_hint, (torch.nn.Module, AutoModel)):
+        if self.type_hint is not None and not issubclass(self.type_hint, torch.nn.Module):
             kwargs.pop("torch_dtype", None)
 
         if self.type_hint is None:
             try:
+                from diffusers import AutoModel
+
                 component = AutoModel.from_pretrained(pretrained_model_name_or_path, **load_kwargs, **kwargs)
             except Exception as e:
                 raise ValueError(f"Unable to load {self.name} without `type_hint`: {e}")
@@ -332,6 +332,12 @@ class ComponentSpec:
                 else getattr(self.type_hint, "from_pretrained")
             )
 
+            # `torch_dtype` is not an accepted parameter for tokenizers and processors.
+            # As a result, it gets stored in `init_kwargs`, which are written to the config
+            # during save. This causes JSON serialization to fail when saving the component.
+            if not issubclass(self.type_hint, torch.nn.Module):
+                kwargs.pop("torch_dtype", None)
+
             try:
                 component = load_method(pretrained_model_name_or_path, **load_kwargs, **kwargs)
             except Exception as e:

src/diffusers/pipelines/flux2/pipeline_flux2.py

@@ -744,7 +744,7 @@ class Flux2Pipeline(DiffusionPipeline, Flux2LoraLoaderMixin):
     @replace_example_docstring(EXAMPLE_DOC_STRING)
     def __call__(
         self,
-        image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
+        image: list[PIL.Image.Image, PIL.Image.Image] | None = None,
         prompt: str | list[str] = None,
         height: int | None = None,
         width: int | None = None,

tests/models/transformers/test_models_transformer_flux.py

@@ -41,7 +41,6 @@ from ..testing_utils import (
     ModelOptCompileTesterMixin,
     ModelOptTesterMixin,
     ModelTesterMixin,
-    PyramidAttentionBroadcastTesterMixin,
     QuantoCompileTesterMixin,
     QuantoTesterMixin,
     SingleFileTesterMixin,
@@ -219,6 +218,10 @@ class TestFluxTransformerMemory(FluxTransformerTesterConfig, MemoryTesterMixin):
 class TestFluxTransformerTraining(FluxTransformerTesterConfig, TrainingTesterMixin):
     """Training tests for Flux Transformer."""
 
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"FluxTransformer2DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
 
 class TestFluxTransformerAttention(FluxTransformerTesterConfig, AttentionTesterMixin):
     """Attention processor tests for Flux Transformer."""
@@ -412,10 +415,6 @@ class TestFluxTransformerBitsAndBytesCompile(FluxTransformerTesterConfig, BitsAn
     """BitsAndBytes + compile tests for Flux Transformer."""
 
 
-class TestFluxTransformerPABCache(FluxTransformerTesterConfig, PyramidAttentionBroadcastTesterMixin):
-    """PyramidAttentionBroadcast cache tests for Flux Transformer."""
-
-
 class TestFluxTransformerFBCCache(FluxTransformerTesterConfig, FirstBlockCacheTesterMixin):
     """FirstBlockCache tests for Flux Transformer."""
 

tests/models/transformers/test_models_transformer_flux2.py

@@ -13,48 +13,95 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-
 import torch
 
-from diffusers import Flux2Transformer2DModel, attention_backend
+from diffusers import Flux2Transformer2DModel
+from diffusers.models.transformers.transformer_flux2 import (
+    Flux2KVAttnProcessor,
+    Flux2KVCache,
+    Flux2KVLayerCache,
+    Flux2KVParallelSelfAttnProcessor,
+)
+from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..test_modeling_common import LoraHotSwappingForModelTesterMixin, ModelTesterMixin, TorchCompileTesterMixin
+from ..testing_utils import (
+    AttentionTesterMixin,
+    BaseModelTesterConfig,
+    BitsAndBytesTesterMixin,
+    ContextParallelTesterMixin,
+    GGUFCompileTesterMixin,
+    GGUFTesterMixin,
+    LoraHotSwappingForModelTesterMixin,
+    LoraTesterMixin,
+    MemoryTesterMixin,
+    ModelTesterMixin,
+    TorchAoCompileTesterMixin,
+    TorchAoTesterMixin,
+    TorchCompileTesterMixin,
+    TrainingTesterMixin,
+)
 
 
 enable_full_determinism()
 
 
-class Flux2TransformerTests(ModelTesterMixin, unittest.TestCase):
-    model_class = Flux2Transformer2DModel
-    main_input_name = "hidden_states"
-    # We override the items here because the transformer under consideration is small.
-    model_split_percents = [0.7, 0.6, 0.6]
-
-    # Skip setting testing with default: AttnProcessor
-    uses_custom_attn_processor = True
+class Flux2TransformerTesterConfig(BaseModelTesterConfig):
+    @property
+    def model_class(self):
+        return Flux2Transformer2DModel
 
     @property
-    def dummy_input(self):
-        return self.prepare_dummy_input()
-
-    @property
-    def input_shape(self):
+    def output_shape(self) -> tuple[int, int]:
         return (16, 4)
 
     @property
-    def output_shape(self):
+    def input_shape(self) -> tuple[int, int]:
         return (16, 4)
 
-    def prepare_dummy_input(self, height=4, width=4):
+    @property
+    def model_split_percents(self) -> list:
+        # We override the items here because the transformer under consideration is small.
+        return [0.7, 0.6, 0.6]
+
+    @property
+    def main_input_name(self) -> str:
+        return "hidden_states"
+
+    @property
+    def uses_custom_attn_processor(self) -> bool:
+        # Skip setting testing with default: AttnProcessor
+        return True
+
+    @property
+    def generator(self):
+        return torch.Generator("cpu").manual_seed(0)
+
+    def get_init_dict(self) -> dict[str, int | list[int]]:
+        return {
+            "patch_size": 1,
+            "in_channels": 4,
+            "num_layers": 1,
+            "num_single_layers": 1,
+            "attention_head_dim": 16,
+            "num_attention_heads": 2,
+            "joint_attention_dim": 32,
+            "timestep_guidance_channels": 256,  # Hardcoded in original code
+            "axes_dims_rope": [4, 4, 4, 4],
+        }
+
+    def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:
         batch_size = 1
         num_latent_channels = 4
         sequence_length = 48
         embedding_dim = 32
 
-        hidden_states = torch.randn((batch_size, height * width, num_latent_channels)).to(torch_device)
-        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
+        hidden_states = randn_tensor(
+            (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device
+        )
+        encoder_hidden_states = randn_tensor(
+            (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
+        )
 
         t_coords = torch.arange(1)
         h_coords = torch.arange(height)
@@ -82,8 +129,286 @@ class Flux2TransformerTests(ModelTesterMixin, unittest.TestCase):
             "guidance": guidance,
         }
 
-    def prepare_init_args_and_inputs_for_common(self):
-        init_dict = {
+
+class TestFlux2Transformer(Flux2TransformerTesterConfig, ModelTesterMixin):
+    pass
+
+
+class TestFlux2TransformerMemory(Flux2TransformerTesterConfig, MemoryTesterMixin):
+    """Memory optimization tests for Flux2 Transformer."""
+
+
+class TestFlux2TransformerTraining(Flux2TransformerTesterConfig, TrainingTesterMixin):
+    """Training tests for Flux2 Transformer."""
+
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"Flux2Transformer2DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+
+class TestFlux2TransformerAttention(Flux2TransformerTesterConfig, AttentionTesterMixin):
+    """Attention processor tests for Flux2 Transformer."""
+
+
+class TestFlux2TransformerContextParallel(Flux2TransformerTesterConfig, ContextParallelTesterMixin):
+    """Context Parallel inference tests for Flux2 Transformer."""
+
+
+class TestFlux2TransformerLoRA(Flux2TransformerTesterConfig, LoraTesterMixin):
+    """LoRA adapter tests for Flux2 Transformer."""
+
+
+class TestFlux2TransformerLoRAHotSwap(Flux2TransformerTesterConfig, LoraHotSwappingForModelTesterMixin):
+    """LoRA hot-swapping tests for Flux2 Transformer."""
+
+    @property
+    def different_shapes_for_compilation(self):
+        return [(4, 4), (4, 8), (8, 8)]
+
+    def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:
+        """Override to support dynamic height/width for LoRA hotswap tests."""
+        batch_size = 1
+        num_latent_channels = 4
+        sequence_length = 48
+        embedding_dim = 32
+
+        hidden_states = randn_tensor(
+            (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device
+        )
+        encoder_hidden_states = randn_tensor(
+            (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
+        )
+
+        t_coords = torch.arange(1)
+        h_coords = torch.arange(height)
+        w_coords = torch.arange(width)
+        l_coords = torch.arange(1)
+        image_ids = torch.cartesian_prod(t_coords, h_coords, w_coords, l_coords)
+        image_ids = image_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        text_t_coords = torch.arange(1)
+        text_h_coords = torch.arange(1)
+        text_w_coords = torch.arange(1)
+        text_l_coords = torch.arange(sequence_length)
+        text_ids = torch.cartesian_prod(text_t_coords, text_h_coords, text_w_coords, text_l_coords)
+        text_ids = text_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+        guidance = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+
+        return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "img_ids": image_ids,
+            "txt_ids": text_ids,
+            "timestep": timestep,
+            "guidance": guidance,
+        }
+
+
+class TestFlux2TransformerCompile(Flux2TransformerTesterConfig, TorchCompileTesterMixin):
+    @property
+    def different_shapes_for_compilation(self):
+        return [(4, 4), (4, 8), (8, 8)]
+
+    def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:
+        """Override to support dynamic height/width for compilation tests."""
+        batch_size = 1
+        num_latent_channels = 4
+        sequence_length = 48
+        embedding_dim = 32
+
+        hidden_states = randn_tensor(
+            (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device
+        )
+        encoder_hidden_states = randn_tensor(
+            (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
+        )
+
+        t_coords = torch.arange(1)
+        h_coords = torch.arange(height)
+        w_coords = torch.arange(width)
+        l_coords = torch.arange(1)
+        image_ids = torch.cartesian_prod(t_coords, h_coords, w_coords, l_coords)
+        image_ids = image_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        text_t_coords = torch.arange(1)
+        text_h_coords = torch.arange(1)
+        text_w_coords = torch.arange(1)
+        text_l_coords = torch.arange(sequence_length)
+        text_ids = torch.cartesian_prod(text_t_coords, text_h_coords, text_w_coords, text_l_coords)
+        text_ids = text_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+        guidance = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+
+        return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "img_ids": image_ids,
+            "txt_ids": text_ids,
+            "timestep": timestep,
+            "guidance": guidance,
+        }
+
+
+class TestFlux2TransformerBitsAndBytes(Flux2TransformerTesterConfig, BitsAndBytesTesterMixin):
+    """BitsAndBytes quantization tests for Flux2 Transformer."""
+
+
+class TestFlux2TransformerTorchAo(Flux2TransformerTesterConfig, TorchAoTesterMixin):
+    """TorchAO quantization tests for Flux2 Transformer."""
+
+
+class TestFlux2TransformerGGUF(Flux2TransformerTesterConfig, GGUFTesterMixin):
+    """GGUF quantization tests for Flux2 Transformer."""
+
+    @property
+    def gguf_filename(self):
+        return "https://huggingface.co/unsloth/FLUX.2-dev-GGUF/blob/main/flux2-dev-Q2_K.gguf"
+
+    @property
+    def torch_dtype(self):
+        return torch.bfloat16
+
+    def get_dummy_inputs(self):
+        """Override to provide inputs matching the real FLUX2 model dimensions.
+
+        Flux2 defaults: in_channels=128, joint_attention_dim=15360
+        """
+        batch_size = 1
+        height = 64
+        width = 64
+        sequence_length = 512
+
+        hidden_states = randn_tensor(
+            (batch_size, height * width, 128), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+        )
+        encoder_hidden_states = randn_tensor(
+            (batch_size, sequence_length, 15360), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+        )
+
+        # Flux2 uses 4D image/text IDs (t, h, w, l)
+        t_coords = torch.arange(1)
+        h_coords = torch.arange(height)
+        w_coords = torch.arange(width)
+        l_coords = torch.arange(1)
+        image_ids = torch.cartesian_prod(t_coords, h_coords, w_coords, l_coords)
+        image_ids = image_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        text_t_coords = torch.arange(1)
+        text_h_coords = torch.arange(1)
+        text_w_coords = torch.arange(1)
+        text_l_coords = torch.arange(sequence_length)
+        text_ids = torch.cartesian_prod(text_t_coords, text_h_coords, text_w_coords, text_l_coords)
+        text_ids = text_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        timestep = torch.tensor([1.0]).to(torch_device, self.torch_dtype)
+        guidance = torch.tensor([3.5]).to(torch_device, self.torch_dtype)
+
+        return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "img_ids": image_ids,
+            "txt_ids": text_ids,
+            "timestep": timestep,
+            "guidance": guidance,
+        }
+
+
+class TestFlux2TransformerTorchAoCompile(Flux2TransformerTesterConfig, TorchAoCompileTesterMixin):
+    """TorchAO + compile tests for Flux2 Transformer."""
+
+
+class TestFlux2TransformerGGUFCompile(Flux2TransformerTesterConfig, GGUFCompileTesterMixin):
+    """GGUF + compile tests for Flux2 Transformer."""
+
+    @property
+    def gguf_filename(self):
+        return "https://huggingface.co/unsloth/FLUX.2-dev-GGUF/blob/main/flux2-dev-Q2_K.gguf"
+
+    @property
+    def torch_dtype(self):
+        return torch.bfloat16
+
+    def get_dummy_inputs(self):
+        """Override to provide inputs matching the real FLUX2 model dimensions.
+
+        Flux2 defaults: in_channels=128, joint_attention_dim=15360
+        """
+        batch_size = 1
+        height = 64
+        width = 64
+        sequence_length = 512
+
+        hidden_states = randn_tensor(
+            (batch_size, height * width, 128), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+        )
+        encoder_hidden_states = randn_tensor(
+            (batch_size, sequence_length, 15360), generator=self.generator, device=torch_device, dtype=self.torch_dtype
+        )
+
+        # Flux2 uses 4D image/text IDs (t, h, w, l)
+        t_coords = torch.arange(1)
+        h_coords = torch.arange(height)
+        w_coords = torch.arange(width)
+        l_coords = torch.arange(1)
+        image_ids = torch.cartesian_prod(t_coords, h_coords, w_coords, l_coords)
+        image_ids = image_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        text_t_coords = torch.arange(1)
+        text_h_coords = torch.arange(1)
+        text_w_coords = torch.arange(1)
+        text_l_coords = torch.arange(sequence_length)
+        text_ids = torch.cartesian_prod(text_t_coords, text_h_coords, text_w_coords, text_l_coords)
+        text_ids = text_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        timestep = torch.tensor([1.0]).to(torch_device, self.torch_dtype)
+        guidance = torch.tensor([3.5]).to(torch_device, self.torch_dtype)
+
+        return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "img_ids": image_ids,
+            "txt_ids": text_ids,
+            "timestep": timestep,
+            "guidance": guidance,
+        }
+
+
+class Flux2TransformerKVCacheTesterConfig(BaseModelTesterConfig):
+    num_ref_tokens = 4
+
+    @property
+    def model_class(self):
+        return Flux2Transformer2DModel
+
+    @property
+    def output_shape(self) -> tuple[int, int]:
+        return (16, 4)
+
+    @property
+    def input_shape(self) -> tuple[int, int]:
+        return (16, 4)
+
+    @property
+    def model_split_percents(self) -> list:
+        return [0.7, 0.6, 0.6]
+
+    @property
+    def main_input_name(self) -> str:
+        return "hidden_states"
+
+    @property
+    def uses_custom_attn_processor(self) -> bool:
+        return True
+
+    @property
+    def generator(self):
+        return torch.Generator("cpu").manual_seed(0)
+
+    def get_init_dict(self) -> dict[str, int | list[int]]:
+        return {
             "patch_size": 1,
             "in_channels": 4,
             "num_layers": 1,
@@ -91,72 +416,210 @@ class Flux2TransformerTests(ModelTesterMixin, unittest.TestCase):
             "attention_head_dim": 16,
             "num_attention_heads": 2,
             "joint_attention_dim": 32,
-            "timestep_guidance_channels": 256,  # Hardcoded in original code
+            "timestep_guidance_channels": 256,
             "axes_dims_rope": [4, 4, 4, 4],
         }
 
-        inputs_dict = self.dummy_input
-        return init_dict, inputs_dict
+    def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]:
+        batch_size = 1
+        num_latent_channels = 4
+        sequence_length = 48
+        embedding_dim = 32
+        num_ref_tokens = self.num_ref_tokens
 
-    # TODO (Daniel, Sayak): We can remove this test.
-    def test_flux2_consistency(self, seed=0):
-        torch.manual_seed(seed)
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        ref_hidden_states = randn_tensor(
+            (batch_size, num_ref_tokens, num_latent_channels), generator=self.generator, device=torch_device
+        )
+        img_hidden_states = randn_tensor(
+            (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device
+        )
+        hidden_states = torch.cat([ref_hidden_states, img_hidden_states], dim=1)
 
-        torch.manual_seed(seed)
-        model = self.model_class(**init_dict)
-        # state_dict = model.state_dict()
-        # for key, param in state_dict.items():
-        #     print(f"{key} | {param.shape}")
-        # torch.save(state_dict, "/raid/daniel_gu/test_flux2_params/diffusers.pt")
+        encoder_hidden_states = randn_tensor(
+            (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
+        )
+
+        ref_t_coords = torch.arange(1)
+        ref_h_coords = torch.arange(num_ref_tokens)
+        ref_w_coords = torch.arange(1)
+        ref_l_coords = torch.arange(1)
+        ref_ids = torch.cartesian_prod(ref_t_coords, ref_h_coords, ref_w_coords, ref_l_coords)
+        ref_ids = ref_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        t_coords = torch.arange(1)
+        h_coords = torch.arange(height)
+        w_coords = torch.arange(width)
+        l_coords = torch.arange(1)
+        image_ids = torch.cartesian_prod(t_coords, h_coords, w_coords, l_coords)
+        image_ids = image_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+        image_ids = torch.cat([ref_ids, image_ids], dim=1)
+
+        text_t_coords = torch.arange(1)
+        text_h_coords = torch.arange(1)
+        text_w_coords = torch.arange(1)
+        text_l_coords = torch.arange(sequence_length)
+        text_ids = torch.cartesian_prod(text_t_coords, text_h_coords, text_w_coords, text_l_coords)
+        text_ids = text_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+
+        timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+        guidance = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+
+        return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "img_ids": image_ids,
+            "txt_ids": text_ids,
+            "timestep": timestep,
+            "guidance": guidance,
+        }
+
+
+class TestFlux2TransformerKVCache(Flux2TransformerKVCacheTesterConfig):
+    """KV cache tests for Flux2 Transformer."""
+
+    def test_kv_layer_cache_store_and_get(self):
+        cache = Flux2KVLayerCache()
+        k = torch.randn(1, 4, 2, 16)
+        v = torch.randn(1, 4, 2, 16)
+        cache.store(k, v)
+        k_out, v_out = cache.get()
+        assert torch.equal(k, k_out)
+        assert torch.equal(v, v_out)
+
+    def test_kv_layer_cache_get_before_store_raises(self):
+        cache = Flux2KVLayerCache()
+        try:
+            cache.get()
+            assert False, "Expected RuntimeError"
+        except RuntimeError:
+            pass
+
+    def test_kv_layer_cache_clear(self):
+        cache = Flux2KVLayerCache()
+        cache.store(torch.randn(1, 4, 2, 16), torch.randn(1, 4, 2, 16))
+        cache.clear()
+        assert cache.k_ref is None
+        assert cache.v_ref is None
+
+    def test_kv_cache_structure(self):
+        num_double = 3
+        num_single = 2
+        cache = Flux2KVCache(num_double, num_single)
+        assert len(cache.double_block_caches) == num_double
+        assert len(cache.single_block_caches) == num_single
+        assert cache.num_ref_tokens == 0
+
+        for i in range(num_double):
+            assert isinstance(cache.get_double(i), Flux2KVLayerCache)
+        for i in range(num_single):
+            assert isinstance(cache.get_single(i), Flux2KVLayerCache)
+
+    def test_kv_cache_clear(self):
+        cache = Flux2KVCache(2, 1)
+        cache.num_ref_tokens = 4
+        cache.get_double(0).store(torch.randn(1, 4, 2, 16), torch.randn(1, 4, 2, 16))
+        cache.clear()
+        assert cache.num_ref_tokens == 0
+        assert cache.get_double(0).k_ref is None
+
+    def _set_kv_attn_processors(self, model):
+        for block in model.transformer_blocks:
+            block.attn.set_processor(Flux2KVAttnProcessor())
+        for block in model.single_transformer_blocks:
+            block.attn.set_processor(Flux2KVParallelSelfAttnProcessor())
+
+    @torch.no_grad()
+    def test_extract_mode_returns_cache(self):
+        model = self.model_class(**self.get_init_dict())
+        model.to(torch_device)
+        model.eval()
+        self._set_kv_attn_processors(model)
+
+        output = model(
+            **self.get_dummy_inputs(),
+            kv_cache_mode="extract",
+            num_ref_tokens=self.num_ref_tokens,
+            ref_fixed_timestep=0.0,
+        )
+
+        assert output.kv_cache is not None
+        assert isinstance(output.kv_cache, Flux2KVCache)
+        assert output.kv_cache.num_ref_tokens == self.num_ref_tokens
+
+        for layer_cache in output.kv_cache.double_block_caches:
+            assert layer_cache.k_ref is not None
+            assert layer_cache.v_ref is not None
+
+        for layer_cache in output.kv_cache.single_block_caches:
+            assert layer_cache.k_ref is not None
+            assert layer_cache.v_ref is not None
+
+    @torch.no_grad()
+    def test_extract_mode_output_shape(self):
+        model = self.model_class(**self.get_init_dict())
         model.to(torch_device)
         model.eval()
 
-        with attention_backend("native"):
-            with torch.no_grad():
-                output = model(**inputs_dict)
+        height, width = 4, 4
+        output = model(
+            **self.get_dummy_inputs(height=height, width=width),
+            kv_cache_mode="extract",
+            num_ref_tokens=self.num_ref_tokens,
+            ref_fixed_timestep=0.0,
+        )
 
-                if isinstance(output, dict):
-                    output = output.to_tuple()[0]
+        assert output.sample.shape == (1, height * width, 4)
 
-        self.assertIsNotNone(output)
+    @torch.no_grad()
+    def test_cached_mode_uses_cache(self):
+        model = self.model_class(**self.get_init_dict())
+        model.to(torch_device)
+        model.eval()
 
-        # input & output have to have the same shape
-        input_tensor = inputs_dict[self.main_input_name]
-        expected_shape = input_tensor.shape
-        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
+        height, width = 4, 4
+        extract_output = model(
+            **self.get_dummy_inputs(height=height, width=width),
+            kv_cache_mode="extract",
+            num_ref_tokens=self.num_ref_tokens,
+            ref_fixed_timestep=0.0,
+        )
 
-        # Check against expected slice
-        # fmt: off
-        expected_slice = torch.tensor([-0.3662, 0.4844, 0.6334, -0.3497, 0.2162, 0.0188, 0.0521, -0.2061, -0.2041, -0.0342, -0.7107, 0.4797, -0.3280, 0.7059, -0.0849, 0.4416])
-        # fmt: on
+        base_config = Flux2TransformerTesterConfig()
+        cached_inputs = base_config.get_dummy_inputs(height=height, width=width)
+        cached_output = model(
+            **cached_inputs,
+            kv_cache=extract_output.kv_cache,
+            kv_cache_mode="cached",
+        )
 
-        flat_output = output.cpu().flatten()
-        generated_slice = torch.cat([flat_output[:8], flat_output[-8:]])
-        self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-4))
+        assert cached_output.sample.shape == (1, height * width, 4)
+        assert cached_output.kv_cache is None
 
-    def test_gradient_checkpointing_is_applied(self):
-        expected_set = {"Flux2Transformer2DModel"}
-        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+    @torch.no_grad()
+    def test_extract_return_dict_false(self):
+        model = self.model_class(**self.get_init_dict())
+        model.to(torch_device)
+        model.eval()
 
+        output = model(
+            **self.get_dummy_inputs(),
+            kv_cache_mode="extract",
+            num_ref_tokens=self.num_ref_tokens,
+            ref_fixed_timestep=0.0,
+            return_dict=False,
+        )
 
-class Flux2TransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
-    model_class = Flux2Transformer2DModel
-    different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]
+        assert isinstance(output, tuple)
+        assert len(output) == 2
+        assert isinstance(output[1], Flux2KVCache)
 
-    def prepare_init_args_and_inputs_for_common(self):
-        return Flux2TransformerTests().prepare_init_args_and_inputs_for_common()
+    @torch.no_grad()
+    def test_no_kv_cache_mode_returns_no_cache(self):
+        model = self.model_class(**self.get_init_dict())
+        model.to(torch_device)
+        model.eval()
 
-    def prepare_dummy_input(self, height, width):
-        return Flux2TransformerTests().prepare_dummy_input(height=height, width=width)
+        base_config = Flux2TransformerTesterConfig()
+        output = model(**base_config.get_dummy_inputs())
 
-
-class Flux2TransformerLoRAHotSwapTests(LoraHotSwappingForModelTesterMixin, unittest.TestCase):
-    model_class = Flux2Transformer2DModel
-    different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)]
-
-    def prepare_init_args_and_inputs_for_common(self):
-        return Flux2TransformerTests().prepare_init_args_and_inputs_for_common()
-
-    def prepare_dummy_input(self, height, width):
-        return Flux2TransformerTests().prepare_dummy_input(height=height, width=width)
+        assert output.kv_cache is None