diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index 9cbe5bc8108f..38dba711cb2a 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -175,6 +175,7 @@ def _register_transformer_blocks_metadata(): from ..models.transformers.transformer_bria import BriaTransformerBlock from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock from ..models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock + from ..models.transformers.transformer_flux2 import Flux2SingleTransformerBlock, Flux2TransformerBlock from ..models.transformers.transformer_hunyuan_video import ( HunyuanVideoSingleTransformerBlock, HunyuanVideoTokenReplaceSingleTransformerBlock, @@ -246,6 +247,22 @@ def _register_transformer_blocks_metadata(): ), ) + # Flux2 + TransformerBlockRegistry.register( + model_class=Flux2TransformerBlock, + metadata=TransformerBlockMetadata( + return_hidden_states_index=1, + return_encoder_hidden_states_index=0, + ), + ) + TransformerBlockRegistry.register( + model_class=Flux2SingleTransformerBlock, + metadata=TransformerBlockMetadata( + return_hidden_states_index=1, + return_encoder_hidden_states_index=0, + ), + ) + # HunyuanVideo TransformerBlockRegistry.register( model_class=HunyuanVideoTransformerBlock, diff --git a/src/diffusers/models/transformers/transformer_flux2.py b/src/diffusers/models/transformers/transformer_flux2.py index 17c8bd0ffd52..6028f475474e 100644 --- a/src/diffusers/models/transformers/transformer_flux2.py +++ b/src/diffusers/models/transformers/transformer_flux2.py @@ -817,18 +817,13 @@ def __init__( def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | None, + encoder_hidden_states: torch.Tensor, temb_mod: torch.Tensor, image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, joint_attention_kwargs: dict[str, Any] | None = None, - split_hidden_states: bool = False, - text_seq_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - # If encoder_hidden_states is None, hidden_states is assumed to have encoder_hidden_states already - # concatenated - if encoder_hidden_states is not None: - text_seq_len = encoder_hidden_states.shape[1] - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) mod_shift, mod_scale, mod_gate = Flux2Modulation.split(temb_mod, 1)[0] @@ -846,11 +841,8 @@ def forward( if hidden_states.dtype == torch.float16: hidden_states = hidden_states.clip(-65504, 65504) - if split_hidden_states: - encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:] - return encoder_hidden_states, hidden_states - else: - return hidden_states + encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:] + return encoder_hidden_states, hidden_states class Flux2TransformerBlock(nn.Module): @@ -1326,12 +1318,9 @@ def forward( joint_attention_kwargs=kv_attn_kwargs, ) - # Concatenate text and image streams for single-block inference - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) - # Blend single block modulation for extract mode: [txt_mod, ref_mod, img_mod] if kv_cache_mode == "extract" and num_ref_tokens > 0: - total_single_len = hidden_states.shape[1] + total_single_len = num_txt_tokens + hidden_states.shape[1] single_stream_mod = _blend_single_block_mods( single_stream_mod, ref_single_mod, num_txt_tokens, num_ref_tokens, total_single_len ) @@ -1348,28 +1337,26 @@ def forward( kv_attn_kwargs_single["kv_cache"] = kv_cache.get_single(index_block) if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( + encoder_hidden_states, hidden_states = self._gradient_checkpointing_func( block, hidden_states, - None, + encoder_hidden_states, single_stream_mod, concat_rotary_emb, kv_attn_kwargs_single, ) else: - hidden_states = block( + encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, - encoder_hidden_states=None, + encoder_hidden_states=encoder_hidden_states, temb_mod=single_stream_mod, image_rotary_emb=concat_rotary_emb, joint_attention_kwargs=kv_attn_kwargs_single, ) - # Remove text tokens (and ref tokens in extract mode) from concatenated stream + # Remove ref tokens (extract mode only) from the image stream if kv_cache_mode == "extract" and num_ref_tokens > 0: - hidden_states = hidden_states[:, num_txt_tokens + num_ref_tokens :, ...] - else: - hidden_states = hidden_states[:, num_txt_tokens:, ...] + hidden_states = hidden_states[:, num_ref_tokens:, ...] # 7. Output layers hidden_states = self.norm_out(hidden_states, temb) diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index 9546fdb5d969..2aa4532ed8d3 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -30,6 +30,7 @@ BaseModelTesterConfig, BitsAndBytesTesterMixin, ContextParallelTesterMixin, + FirstBlockCacheTesterMixin, GGUFCompileTesterMixin, GGUFTesterMixin, LoraHotSwappingForModelTesterMixin, @@ -153,6 +154,10 @@ class TestFlux2TransformerContextParallel(Flux2TransformerTesterConfig, ContextP """Context Parallel inference tests for Flux2 Transformer.""" +class TestFlux2TransformerFBCCache(Flux2TransformerTesterConfig, FirstBlockCacheTesterMixin): + """FirstBlockCache tests for Flux2 Transformer.""" + + class TestFlux2TransformerLoRA(Flux2TransformerTesterConfig, LoraTesterMixin): """LoRA adapter tests for Flux2 Transformer.""" diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein.py b/tests/pipelines/flux2/test_pipeline_flux2_klein.py index a38eb4de2a90..056f493eec0f 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein.py +++ b/tests/pipelines/flux2/test_pipeline_flux2_klein.py @@ -23,6 +23,7 @@ ) from ..testing_utils import ( BasePipelineTesterConfig, + FirstBlockCacheTesterMixin, MemoryTesterMixin, PipelineTesterMixin, check_qkv_fused_layers_exist, @@ -197,6 +198,10 @@ class TestFlux2KleinPipelineMemory(Flux2KleinPipelineTesterConfig, MemoryTesterM """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux2 Klein pipeline.""" +class TestFlux2KleinPipelineFirstBlockCache(Flux2KleinPipelineTesterConfig, FirstBlockCacheTesterMixin): + """First Block Cache tests for the Flux2 Klein pipeline.""" + + @require_torch_neuron class TestFlux2KleinPipelineIntegration: ckpt_id = "black-forest-labs/FLUX.2-klein-4B"