# Copyright 2025 Black Forest Labs, The HuggingFace Team and The InstantX Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn as nn from diffusers.models.attention import AttentionModuleMixin from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.normalization import AdaLayerNormContinuous from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.runtime.distributed import ( divide, get_tp_world_size, ) from sglang.multimodal_gen.runtime.distributed.sp_shard_utils import ( build_shard_plan, join_seqs, shard_like, shard_seq_prefix, should_shard_text, split_seqs, tail_attn_meta, ) from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import ( RMSNorm, apply_qk_norm_with_optional_rope, ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( ModelOptFp4Config, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name def _get_qkv_projections( attn: "Flux2Attention", hidden_states, encoder_hidden_states=None ): if attn.use_fused_qkv: qkv, _ = attn.to_qkv(hidden_states) query, key, value = [t.contiguous() for t in qkv.chunk(3, dim=-1)] else: query, _ = attn.to_q(hidden_states) key, _ = attn.to_k(hidden_states) value, _ = attn.to_v(hidden_states) encoder_query = encoder_key = encoder_value = None if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None: if attn.use_fused_added_qkv: added_qkv, _ = attn.to_added_qkv(encoder_hidden_states) encoder_query, encoder_key, encoder_value = [ t.contiguous() for t in added_qkv.chunk(3, dim=-1) ] else: encoder_query, _ = attn.add_q_proj(encoder_hidden_states) encoder_key, _ = attn.add_k_proj(encoder_hidden_states) encoder_value, _ = attn.add_v_proj(encoder_hidden_states) return query, key, value, encoder_query, encoder_key, encoder_value class Flux2SwiGLU(nn.Module): """ Flux 2 uses a SwiGLU-style activation in the transformer feedforward sub-blocks, but with the linear projection layer fused into the first linear layer of the FF sub-block. Thus, this module has no trainable parameters. """ def __init__(self): super().__init__() self.gate_fn = nn.SiLU() def forward(self, x: torch.Tensor) -> torch.Tensor: x1, x2 = x.chunk(2, dim=-1) x = self.gate_fn(x1) * x2 return x class Flux2FeedForward(nn.Module): def __init__( self, dim: int, dim_out: Optional[int] = None, mult: float = 3.0, inner_dim: Optional[int] = None, bias: bool = False, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() if inner_dim is None: inner_dim = int(dim * mult) dim_out = dim_out or dim # Flux2SwiGLU will reduce the dimension by half self.linear_in = MergedColumnParallelLinear( dim, [inner_dim, inner_dim], bias=bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.linear_in" if prefix else "linear_in", ) self.act_fn = Flux2SwiGLU() self.linear_out = RowParallelLinear( inner_dim, dim_out, bias=bias, input_is_parallel=True, quant_config=quant_config, prefix=f"{prefix}.linear_out" if prefix else "linear_out", ) def forward(self, x: torch.Tensor) -> torch.Tensor: x, _ = self.linear_in(x) x = self.act_fn(x) x, _ = self.linear_out(x) return x class Flux2Attention(torch.nn.Module, AttentionModuleMixin): def __init__( self, query_dim: int, num_heads: int = 8, dim_head: int = 64, dropout: float = 0.0, bias: bool = False, added_kv_proj_dim: Optional[int] = None, added_proj_bias: Optional[bool] = True, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, elementwise_affine: bool = True, supported_attention_backends: set[AttentionBackendEnum] | None = None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * num_heads self.query_dim = query_dim self.out_dim = out_dim if out_dim is not None else query_dim self.heads = out_dim // dim_head if out_dim is not None else num_heads self.tp_size = get_tp_world_size() self.local_heads = divide(self.heads, self.tp_size) self.local_inner_dim = divide(self.inner_dim, self.tp_size) self.use_bias = bias self.dropout = dropout self.added_kv_proj_dim = added_kv_proj_dim self.added_proj_bias = added_proj_bias # Some FLUX.2 NVFP4 checkpoints store Q/K/V packed as a single tensor, while # ModelOpt's standard diffusers export keeps the original to_q/to_k/to_v layout. # Only enable the fused loader path for the packed checkpoint family. self.use_fused_qkv = isinstance(quant_config, ModelOptFp4Config) and getattr( quant_config, "checkpoint_uses_packed_qkv", False ) self.use_fused_added_qkv = self.use_fused_qkv if self.use_fused_qkv: self.to_qkv = MergedColumnParallelLinear( query_dim, [self.inner_dim] * 3, bias=bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.to_qkv" if prefix else "to_qkv", ) else: self.to_q = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.to_q" if prefix else "to_q", ) self.to_k = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.to_k" if prefix else "to_k", ) self.to_v = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.to_v" if prefix else "to_v", ) # QK Norm self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps) self.to_out = torch.nn.ModuleList([]) self.to_out.append( RowParallelLinear( self.inner_dim, self.out_dim, bias=out_bias, input_is_parallel=True, quant_config=quant_config, prefix=f"{prefix}.to_out.0" if prefix else "to_out.0", ) ) self.to_out.append(torch.nn.Dropout(dropout)) if added_kv_proj_dim is not None: self.norm_added_q = RMSNorm(dim_head, eps=eps) self.norm_added_k = RMSNorm(dim_head, eps=eps) if self.use_fused_added_qkv: # txt_attn.qkv is always BF16 in the NVFP4 checkpoint — no quant needed self.to_added_qkv = MergedColumnParallelLinear( added_kv_proj_dim, [self.inner_dim] * 3, bias=added_proj_bias, gather_output=False, quant_config=None, prefix=f"{prefix}.to_added_qkv" if prefix else "to_added_qkv", ) else: self.add_q_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.add_q_proj" if prefix else "add_q_proj", ) self.add_k_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.add_k_proj" if prefix else "add_k_proj", ) self.add_v_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.add_v_proj" if prefix else "add_v_proj", ) self.to_add_out = RowParallelLinear( self.inner_dim, query_dim, bias=out_bias, input_is_parallel=True, quant_config=quant_config, prefix=f"{prefix}.to_add_out" if prefix else "to_add_out", ) self.attn = USPAttention( num_heads=self.local_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends=supported_attention_backends, ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, num_replicated_prefix: int = 0, attn_mask: Optional[torch.Tensor] = None, attn_mask_meta: Optional[Dict[str, int]] = None, ) -> torch.Tensor: ( query, key, value, encoder_query, encoder_key, encoder_value, ) = _get_qkv_projections(self, hidden_states, encoder_hidden_states) query = query.unflatten(-1, (self.local_heads, -1)) key = key.unflatten(-1, (self.local_heads, -1)) value = value.unflatten(-1, (self.local_heads, -1)) cos_sin_cache = None if freqs_cis is not None: cos, sin = freqs_cis cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) if self.added_kv_proj_dim is not None: encoder_query = encoder_query.unflatten(-1, (self.local_heads, -1)) encoder_key = encoder_key.unflatten(-1, (self.local_heads, -1)) encoder_value = encoder_value.unflatten(-1, (self.local_heads, -1)) text_seq_len = encoder_query.shape[1] encoder_query, encoder_key = apply_qk_norm_with_optional_rope( q=encoder_query, k=encoder_key, q_norm=self.norm_added_q, k_norm=self.norm_added_k, head_dim=self.head_dim, cos_sin_cache=cos_sin_cache, is_neox=False, allow_inplace=True, ) query, key = apply_qk_norm_with_optional_rope( q=query, k=key, q_norm=self.norm_q, k_norm=self.norm_k, head_dim=self.head_dim, cos_sin_cache=cos_sin_cache, is_neox=False, position_offset=text_seq_len, allow_inplace=True, ) # join_seqs relocates any SP text tail-pad behind the image (see # sp_shard.join_seqs for why). sp_txt_pad = (attn_mask_meta or {}).get("local_pad", 0) query = join_seqs(encoder_query, query, sp_txt_pad) key = join_seqs(encoder_key, key, sp_txt_pad) value = join_seqs(encoder_value, value, sp_txt_pad) else: query, key = apply_qk_norm_with_optional_rope( q=query, k=key, q_norm=self.norm_q, k_norm=self.norm_k, head_dim=self.head_dim, cos_sin_cache=cos_sin_cache, is_neox=False, allow_inplace=True, ) hidden_states = self.attn( query, key, value, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta, num_replicated_prefix=num_replicated_prefix, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) if encoder_hidden_states is not None: encoder_hidden_states, hidden_states = split_seqs( hidden_states, encoder_hidden_states.shape[1], sp_txt_pad ) encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states) hidden_states, _ = self.to_out[0](hidden_states) hidden_states = self.to_out[1](hidden_states) if encoder_hidden_states is not None: return hidden_states, encoder_hidden_states else: return hidden_states class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): """ Flux 2 parallel self-attention for the Flux 2 single-stream transformer blocks. This implements a parallel transformer block, where the attention QKV projections are fused to the feedforward (FF) input projections, and the attention output projections are fused to the FF output projections. See the [ViT-22B paper](https://arxiv.org/abs/2302.05442) for a visual depiction of this type of transformer block. """ # Does not support QKV fusion as the QKV projections are always fused _supports_qkv_fusion = False def __init__( self, query_dim: int, num_heads: int = 8, dim_head: int = 64, dropout: float = 0.0, bias: bool = False, out_bias: bool = True, eps: float = 1e-5, out_dim: int = None, elementwise_affine: bool = True, mlp_ratio: float = 4.0, mlp_mult_factor: int = 2, supported_attention_backends: set[AttentionBackendEnum] | None = None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * num_heads self.query_dim = query_dim self.out_dim = out_dim if out_dim is not None else query_dim self.heads = out_dim // dim_head if out_dim is not None else num_heads self.tp_size = get_tp_world_size() self.local_heads = divide(self.heads, self.tp_size) self.local_inner_dim = divide(self.inner_dim, self.tp_size) self.use_bias = bias self.dropout = dropout self.mlp_ratio = mlp_ratio self.mlp_hidden_dim = int(query_dim * self.mlp_ratio) self.local_mlp_hidden_dim = divide(self.mlp_hidden_dim, self.tp_size) self.mlp_mult_factor = mlp_mult_factor # Fused QKV projections + MLP input projection self.to_qkv_mlp_proj = MergedColumnParallelLinear( self.query_dim, [self.inner_dim, self.inner_dim, self.inner_dim] + [self.mlp_hidden_dim] * self.mlp_mult_factor, bias=bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.to_qkv_mlp_proj" if prefix else "to_qkv_mlp_proj", ) self.mlp_act_fn = Flux2SwiGLU() # QK Norm self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps) # Fused attention output + MLP output projection. # Input is [attn_shard | mlp_shard] (independently sharded by # MergedColumnParallelLinear), so patch weight loader to pick the # correct non-contiguous columns per rank. self.to_out = RowParallelLinear( self.inner_dim + self.mlp_hidden_dim, self.out_dim, bias=out_bias, input_is_parallel=True, quant_config=quant_config, prefix=f"{prefix}.to_out" if prefix else "to_out", ) if self.tp_size > 1: self._patch_to_out_weight_loader() self.attn = USPAttention( num_heads=self.local_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends=supported_attention_backends, ) def _patch_to_out_weight_loader(self) -> None: inner_dim, mlp_dim = self.inner_dim, self.mlp_hidden_dim tp_size, tp_rank = self.tp_size, self.to_out.tp_rank def _loader(param, loaded_weight): input_dim = getattr(param, "input_dim", None) if input_dim is not None: a = inner_dim // tp_size m = mlp_dim // tp_size attn_cols = loaded_weight.narrow(input_dim, tp_rank * a, a) mlp_cols = loaded_weight.narrow(input_dim, inner_dim + tp_rank * m, m) param.data.copy_(torch.cat([attn_cols, mlp_cols], dim=input_dim)) else: param.data.copy_(loaded_weight) self.to_out.weight_loader = _loader if hasattr(self.to_out.weight, "_weight_loader"): self.to_out.weight._weight_loader = _loader else: self.to_out.weight.weight_loader = _loader def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, num_replicated_prefix: int = 0, **kwargs, ) -> torch.Tensor: attn_mask = kwargs.get("attn_mask") attn_mask_meta = kwargs.get("attn_mask_meta") if attn_mask is None: attn_mask = attention_mask # Parallel in (QKV + MLP in) projection hidden_states, _ = self.to_qkv_mlp_proj(hidden_states) qkv, mlp_hidden_states = torch.split( hidden_states, [ 3 * self.local_inner_dim, self.local_mlp_hidden_dim * self.mlp_mult_factor, ], dim=-1, ) # Handle the attention logic query, key, value = qkv.chunk(3, dim=-1) query = query.unflatten(-1, (self.local_heads, -1)) key = key.unflatten(-1, (self.local_heads, -1)) value = value.unflatten(-1, (self.local_heads, -1)) query = self.norm_q(query) key = self.norm_k(key) if freqs_cis is not None: cos, sin = freqs_cis cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) hidden_states = self.attn( query, key, value, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta, num_replicated_prefix=num_replicated_prefix, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) # Handle the feedforward (FF) logic mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states) # Concatenate and parallel output projection hidden_states = torch.cat([hidden_states, mlp_hidden_states], dim=-1) hidden_states, _ = self.to_out(hidden_states) return hidden_states class Flux2SingleTransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 3.0, eps: float = 1e-6, bias: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) # Note that the MLP in/out linear layers are fused with the attention QKV/out projections, respectively; this # is often called a "parallel" transformer block. See the [ViT-22B paper](https://arxiv.org/abs/2302.05442) # for a visual depiction of this type of transformer block. self.attn = Flux2ParallelSelfAttention( query_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, bias=bias, out_bias=bias, eps=eps, mlp_ratio=mlp_ratio, mlp_mult_factor=2, supported_attention_backends=supported_attention_backends, quant_config=quant_config, prefix=f"{prefix}.attn" if prefix else "attn", ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor], temb_mod_params: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, split_hidden_states: bool = False, text_seq_len: Optional[int] = None, num_replicated_prefix: int = 0, ) -> 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) mod_shift, mod_scale, mod_gate = temb_mod_params norm_hidden_states = self.norm(hidden_states) norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift joint_attention_kwargs = joint_attention_kwargs or {} attn_output = self.attn( hidden_states=norm_hidden_states, freqs_cis=freqs_cis, num_replicated_prefix=num_replicated_prefix, **joint_attention_kwargs, ) hidden_states = hidden_states + mod_gate * attn_output 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 class Flux2TransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 3.0, eps: float = 1e-6, bias: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.attn = Flux2Attention( query_dim=dim, added_kv_proj_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, bias=bias, added_proj_bias=bias, out_bias=bias, eps=eps, supported_attention_backends=supported_attention_backends, quant_config=quant_config, prefix=f"{prefix}.attn" if prefix else "attn", ) self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.ff = Flux2FeedForward( dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias, quant_config=quant_config, prefix=f"{prefix}.ff" if prefix else "ff", ) self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.ff_context = Flux2FeedForward( dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias, quant_config=quant_config, prefix=f"{prefix}.ff_context" if prefix else "ff_context", ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb_mod_params_img: Tuple[ Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ... ], temb_mod_params_txt: Tuple[ Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ... ], freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, num_replicated_prefix: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor]: joint_attention_kwargs = joint_attention_kwargs or {} # Modulation parameters shape: [1, 1, self.dim] (shift_msa, scale_msa, gate_msa), ( shift_mlp, scale_mlp, gate_mlp, ) = temb_mod_params_img (c_shift_msa, c_scale_msa, c_gate_msa), ( c_shift_mlp, c_scale_mlp, c_gate_mlp, ) = temb_mod_params_txt # Img stream norm_hidden_states = self.norm1(hidden_states) norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa # Conditioning txt stream norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) norm_encoder_hidden_states = ( 1 + c_scale_msa ) * norm_encoder_hidden_states + c_shift_msa # Attention on concatenated img + txt stream attention_outputs = self.attn( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, freqs_cis=freqs_cis, num_replicated_prefix=num_replicated_prefix, **joint_attention_kwargs, ) attn_output, context_attn_output = attention_outputs # Process attention outputs for the image stream (`hidden_states`). attn_output = gate_msa * attn_output hidden_states = hidden_states + attn_output norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp ff_output = self.ff(norm_hidden_states) hidden_states = hidden_states + gate_mlp * ff_output # Process attention outputs for the text stream (`encoder_hidden_states`). context_attn_output = c_gate_msa * context_attn_output encoder_hidden_states = encoder_hidden_states + context_attn_output norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) norm_encoder_hidden_states = ( norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp ) context_ff_output = self.ff_context(norm_encoder_hidden_states) encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output if encoder_hidden_states.dtype == torch.float16: encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) return encoder_hidden_states, hidden_states class Flux2TimestepGuidanceEmbeddings(nn.Module): def __init__( self, in_channels: int = 256, embedding_dim: int = 6144, bias: bool = False, guidance_embeds: bool = True, ): super().__init__() self.time_proj = Timesteps( num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0 ) self.timestep_embedder = TimestepEmbedding( in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias ) if guidance_embeds: self.guidance_embedder = TimestepEmbedding( in_channels=in_channels, time_embed_dim=embedding_dim, sample_proj_bias=bias, ) else: self.guidance_embedder = None def forward( self, timestep: torch.Tensor, guidance: Optional[torch.Tensor] = None ) -> torch.Tensor: timesteps_proj = self.time_proj(timestep) timesteps_emb = self.timestep_embedder( timesteps_proj.to(timestep.dtype) ) # (N, D) if guidance is not None and self.guidance_embedder is not None: guidance_proj = self.time_proj(guidance) guidance_emb = self.guidance_embedder( guidance_proj.to(guidance.dtype) ) # (N, D) time_guidance_emb = timesteps_emb + guidance_emb return time_guidance_emb else: return timesteps_emb class Flux2Modulation(nn.Module): def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False): super().__init__() self.mod_param_sets = mod_param_sets self.linear = ColumnParallelLinear( dim, dim * 3 * self.mod_param_sets, bias=bias, gather_output=True ) self.act_fn = nn.SiLU() def forward( self, temb: torch.Tensor ) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: mod = self.act_fn(temb) mod, _ = self.linear(mod) if mod.ndim == 2: mod = mod.unsqueeze(1) mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1) # Return tuple of 3-tuples of modulation params shift/scale/gate return tuple( mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets) ) class Flux2PosEmbed(nn.Module): def __init__(self, theta: int, axes_dim: List[int]): super().__init__() self.rope = NDRotaryEmbedding( rope_dim_list=axes_dim, rope_theta=theta, use_real=False, repeat_interleave_real=False, dtype=( torch.float64 if ( current_platform.is_float64_supported() if hasattr(current_platform, "is_float64_supported") else True ) else torch.float32 ), ) def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: pos = ids.float() # TODO: potential error: flux use n_axes = ids.shape[-1] # see: https://github.com/huggingface/diffusers/blob/17c0e79dbdf53fb6705e9c09cc1a854b84c39249/src/diffusers/models/transformers/transformer_flux.py#L509 freqs_cos, freqs_sin = self.rope.forward_uncached(pos=pos) return freqs_cos.contiguous().float(), freqs_sin.contiguous().float() class Flux2Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): """ The Transformer model introduced in Flux 2. Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ """ param_names_mapping = FluxConfig().arch_config.param_names_mapping scale_shift_swap_params = ("norm_out.linear.weight", "norm_out.linear.bias") # FLUX.2 stays closer to the official diffusers output with Torch SDPA. # The generic FA path still produces a measurable image-level drift here. _supported_attention_backends = { AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.FA, AttentionBackendEnum.AITER, AttentionBackendEnum.AITER_SAGE, } def post_load_weights(self) -> None: if not isinstance(getattr(self, "quant_config", None), ModelOptFp4Config): return # BFL/ComfyUI checkpoints store AdaLN modulation params as [scale, shift], # while diffusers expects [shift, scale]. for param_name in self.scale_shift_swap_params: parts = param_name.split(".") module = self for part in parts[:-1]: module = getattr(module, part) param = getattr(module, parts[-1], None) if param is None: continue half = param.shape[0] // 2 with torch.no_grad(): first_half = param[:half].clone() param[:half] = param[half:] param[half:] = first_half logger.info( "Swapped scale/shift order for %s (BFL → diffusers)", param_name ) def __init__( self, config: FluxConfig, hf_config: dict[str, Any], quant_config: Optional[QuantizationConfig] = None, ): super().__init__(config=config, hf_config=hf_config) patch_size: int = config.patch_size in_channels: int = config.in_channels out_channels: Optional[int] = config.out_channels num_layers: int = config.num_layers num_single_layers: int = config.num_single_layers attention_head_dim: int = config.attention_head_dim num_attention_heads: int = config.num_attention_heads joint_attention_dim: int = config.joint_attention_dim timestep_guidance_channels: int = config.timestep_guidance_channels mlp_ratio: float = config.mlp_ratio axes_dims_rope: Tuple[int, ...] = config.axes_dims_rope rope_theta: int = config.rope_theta eps: float = config.eps guidance_embeds: bool = getattr(config, "guidance_embeds", True) self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim self.guidance_embeds = guidance_embeds quant_config = quant_config if quant_config is not None else config.quant_config self.quant_config = quant_config # 1. Sinusoidal positional embedding for RoPE on image and text tokens self.rotary_emb = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope) # 2. Combined timestep + guidance embedding self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings( in_channels=timestep_guidance_channels, embedding_dim=self.inner_dim, bias=False, guidance_embeds=guidance_embeds, ) # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.) # Two sets of shift/scale/gate modulation parameters for the double stream attn and FF sub-blocks self.double_stream_modulation_img = Flux2Modulation( self.inner_dim, mod_param_sets=2, bias=False ) self.double_stream_modulation_txt = Flux2Modulation( self.inner_dim, mod_param_sets=2, bias=False ) # Only one set of modulation parameters as the attn and FF sub-blocks are run in parallel for single stream self.single_stream_modulation = Flux2Modulation( self.inner_dim, mod_param_sets=1, bias=False ) # 4. Input projections self.x_embedder = ColumnParallelLinear( in_channels, self.inner_dim, bias=False, gather_output=True ) self.context_embedder = ColumnParallelLinear( joint_attention_dim, self.inner_dim, bias=False, gather_output=True ) # 5. Double Stream Transformer Blocks self.transformer_blocks = nn.ModuleList( [ Flux2TransformerBlock( dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, mlp_ratio=mlp_ratio, eps=eps, bias=False, supported_attention_backends=self._supported_attention_backends, quant_config=quant_config, prefix=f"transformer_blocks.{i}", ) for i in range(num_layers) ] ) # 6. Single Stream Transformer Blocks self.single_transformer_blocks = nn.ModuleList( [ Flux2SingleTransformerBlock( dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, mlp_ratio=mlp_ratio, eps=eps, bias=False, supported_attention_backends=self._supported_attention_backends, quant_config=quant_config, prefix=f"single_transformer_blocks.{i}", ) for i in range(num_single_layers) ] ) # 7. Output layers self.norm_out = AdaLayerNormContinuous( self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False, ) self.proj_out = ColumnParallelLinear( self.inner_dim, patch_size * patch_size * self.out_channels, bias=False, gather_output=True, quant_config=quant_config, prefix="proj_out", ) self.layer_names = ["transformer_blocks", "single_transformer_blocks"] def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, timestep: torch.LongTensor = None, guidance: torch.Tensor = None, freqs_cis: torch.Tensor = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> torch.Tensor: """ The [`FluxTransformer2DModel`] forward method. Args: hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): Input `hidden_states`. encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`): Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. timestep ( `torch.LongTensor`): Used to indicate denoising step. joint_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). """ # 0. Handle input arguments if joint_attention_kwargs is not None: joint_attention_kwargs = joint_attention_kwargs.copy() joint_attention_kwargs.pop("scale", 1.0) num_txt_tokens = encoder_hidden_states.shape[1] # 1. Calculate timestep embedding and modulation parameters timestep = timestep.to(hidden_states.dtype) if guidance is not None: guidance = guidance.to(hidden_states.dtype) * 1000 temb = self.time_guidance_embed(timestep, guidance) double_stream_mod_img = self.double_stream_modulation_img(temb) double_stream_mod_txt = self.double_stream_modulation_txt(temb) single_stream_mod = self.single_stream_modulation(temb)[0] # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states) hidden_states, _ = self.x_embedder(hidden_states) encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states) # Shard the replicated text stream across SP ranks (image latents are # already sharded); non-divisible lengths tail-pad the last rank and the # per-request tail meta lets attention skip the pad for free. num_replicated_prefix = num_txt_tokens sp_txt_pad = 0 singles_freqs_cis = freqs_cis if should_shard_text(num_txt_tokens): txt_shard = build_shard_plan(num_txt_tokens) encoder_hidden_states = shard_like(encoder_hidden_states, txt_shard) if freqs_cis is not None: cos, sin = freqs_cis cos = shard_seq_prefix(cos, num_txt_tokens, txt_shard) sin = shard_seq_prefix(sin, num_txt_tokens, txt_shard) freqs_cis = (cos, sin) singles_freqs_cis = freqs_cis num_replicated_prefix = 0 num_txt_tokens = txt_shard.local_len tail_meta = tail_attn_meta( txt_shard, encoder_hidden_states.shape[0], hidden_states.device, image_seq_len=hidden_states.shape[1], ) if tail_meta is not None: joint_attention_kwargs = ( joint_attention_kwargs.copy() if joint_attention_kwargs else {} ) joint_attention_kwargs["attn_mask_meta"] = tail_meta sp_txt_pad = txt_shard.local_pad # The single-stream trunk applies RoPE on the relocated # [txt_real, img, pad] layout; reorder its cache to match. if freqs_cis is not None: t_loc = txt_shard.local_len singles_freqs_cis = ( join_seqs(cos[:t_loc], cos[t_loc:], sp_txt_pad, dim=0), join_seqs(sin[:t_loc], sin[t_loc:], sp_txt_pad, dim=0), ) # 4. Double Stream Transformer Blocks for index_block, block in enumerate(self.transformer_blocks): encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb_mod_params_img=double_stream_mod_img, temb_mod_params_txt=double_stream_mod_txt, freqs_cis=freqs_cis, joint_attention_kwargs=joint_attention_kwargs, num_replicated_prefix=num_replicated_prefix, ) # Concatenate text and image streams for single-block inference; # join_seqs relocates any SP text tail-pad behind the image once for # the whole trunk (see sp_shard.join_seqs for why). txt_real = num_txt_tokens - sp_txt_pad hidden_states = join_seqs(encoder_hidden_states, hidden_states, sp_txt_pad) # 5. Single Stream Transformer Blocks for index_block, block in enumerate(self.single_transformer_blocks): hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=None, temb_mod_params=single_stream_mod, freqs_cis=singles_freqs_cis, joint_attention_kwargs=joint_attention_kwargs, text_seq_len=txt_real, num_replicated_prefix=num_replicated_prefix, ) # Remove text (and any tail pad) from the concatenated stream img_end = hidden_states.shape[1] - sp_txt_pad hidden_states = hidden_states[:, txt_real:img_end, ...] # 6. Output layers hidden_states = self.norm_out(hidden_states, temb) output, _ = self.proj_out(hidden_states) return output EntryClass = Flux2Transformer2DModel