# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # 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, Union import torch import torch.nn as nn from diffusers.models.attention import AttentionModuleMixin from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import ( AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle, ) from torch.nn import LayerNorm as LayerNorm 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.mlp import FeedForward from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( NunchakuConfig, is_nunchaku_available, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import ( CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, ) 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 current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name try: from nunchaku.models.attention import NunchakuFeedForward # type: ignore[import] from nunchaku.models.normalization import ( # type: ignore[import] NunchakuAdaLayerNormZero, NunchakuAdaLayerNormZeroSingle, ) from nunchaku.ops.gemm import ( svdq_gemm_w4a4_cuda as _svdq_gemm_w4a4, # type: ignore[import] ) from nunchaku.ops.quantize import ( svdq_quantize_w4a4_act_fuse_lora_cuda as _svdq_quantize_w4a4, # type: ignore[import] ) _nunchaku_fused_ops_available = True except Exception: NunchakuFeedForward = None NunchakuAdaLayerNormZero = None NunchakuAdaLayerNormZeroSingle = None _svdq_gemm_w4a4 = None _svdq_quantize_w4a4 = None _nunchaku_fused_ops_available = False def _fused_gelu_mlp( x: torch.Tensor, fc1, fc2, pad_size: int = 256, ) -> torch.Tensor: """ Fused GELU MLP matching nunchaku's fused_gelu_mlp kernel path. nunchaku's single-block MLP checkpoint is calibrated for the fused path where: 1. fc1 GEMM + GELU + 0.171875 shift + unsigned re-quantization + fc2.lora_down are all done in a single fused kernel call 2. fc2 GEMM then receives unsigned INT4 activations (act_unsigned=True) Using the sequential path (fc1 → GELU → fc2 with symmetric quantization) is fundamentally incompatible with these wscales, causing visually wrong outputs. """ batch_size, seq_len, channels = x.shape x_2d = x.view(batch_size * seq_len, channels) quantized_x, ascales, lora_act = _svdq_quantize_w4a4( x_2d, lora_down=fc1.proj_down, smooth=fc1.smooth_factor, fp4=fc1.precision == "nvfp4", pad_size=pad_size, ) batch_size_pad = (batch_size * seq_len + pad_size - 1) // pad_size * pad_size is_fp4 = fc2.precision == "nvfp4" qout_act = torch.empty( batch_size_pad, fc1.output_size_per_partition // 2, dtype=torch.uint8, device=x_2d.device, ) if is_fp4: qout_ascales = torch.empty( fc1.output_size_per_partition // 16, batch_size_pad, dtype=torch.float8_e4m3fn, device=x_2d.device, ) else: qout_ascales = torch.empty( fc1.output_size_per_partition // 64, batch_size_pad, dtype=x_2d.dtype, device=x_2d.device, ) qout_lora_act = torch.empty( batch_size_pad, fc2.proj_down.shape[1], dtype=torch.float32, device=x_2d.device ) # fused: fc1 GEMM + GELU + shift + unsigned quantize + fc2.lora_down _svdq_gemm_w4a4( act=quantized_x, wgt=fc1.qweight, qout=qout_act, ascales=ascales, wscales=fc1.wscales, oscales=qout_ascales, lora_act_in=lora_act, lora_up=fc1.proj_up, lora_down=fc2.proj_down, lora_act_out=qout_lora_act, bias=fc1.bias, smooth_factor=fc2.smooth_factor, fp4=is_fp4, alpha=getattr(fc1, "_nunchaku_alpha", None), wcscales=getattr(fc1, "wcscales", None), ) output = torch.empty( batch_size * seq_len, fc2.output_size_per_partition, dtype=x_2d.dtype, device=x_2d.device, ) # fc2 GEMM with unsigned INT4 activations (fused kernel shifted by 0.171875) _svdq_gemm_w4a4( act=qout_act, wgt=fc2.qweight, out=output, ascales=qout_ascales, wscales=fc2.wscales, lora_act_in=qout_lora_act, lora_up=fc2.proj_up, bias=fc2.bias, fp4=is_fp4, alpha=getattr(fc2, "_nunchaku_alpha", None), wcscales=getattr(fc2, "wcscales", None), act_unsigned=True, ) return output.view(batch_size, seq_len, -1) def _get_qkv_projections( attn: "FluxAttention", hidden_states, encoder_hidden_states=None ): if getattr(attn, "use_fused_qkv", False): qkv, _ = attn.to_qkv(hidden_states) query, key, value = [x.contiguous() for x 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 = [ x.contiguous() for x 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 FluxGELU(nn.Module): def __init__( self, dim: int, inner_dim: int, bias: bool = True, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.proj = ColumnParallelLinear( dim, inner_dim, bias=bias, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.proj" if prefix else "proj", ) self.gelu = nn.GELU(approximate="tanh") def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states, _ = self.proj(hidden_states) return self.gelu(hidden_states) class FluxParallelFeedForward(nn.Module): def __init__( self, dim: int, dim_out: Optional[int] = None, mult: int = 4, inner_dim: Optional[int] = None, bias: bool = True, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() if inner_dim is None: inner_dim = int(dim * mult) dim_out = dim_out if dim_out is not None else dim self.net = nn.ModuleList( [ FluxGELU( dim, inner_dim, bias=bias, quant_config=quant_config, prefix=f"{prefix}.net.0" if prefix else "net.0", ), nn.Dropout(0.0), RowParallelLinear( inner_dim, dim_out, bias=bias, input_is_parallel=True, quant_config=quant_config, prefix=f"{prefix}.net.2" if prefix else "net.2", ), ] ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.net[0](hidden_states) hidden_states = self.net[1](hidden_states) hidden_states, _ = self.net[2](hidden_states) return hidden_states class FluxAttention(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, context_pre_only: Optional[bool] = None, pre_only: bool = False, 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.use_bias = bias self.dropout = dropout self.out_dim = out_dim if out_dim is not None else query_dim self.context_pre_only = context_pre_only self.pre_only = pre_only self.heads = out_dim // dim_head if out_dim is not None else num_heads self.tp_size = get_tp_world_size() self.shard_qkv = self.tp_size > 1 and not isinstance( quant_config, NunchakuConfig ) self.local_heads = divide(self.heads, self.tp_size) self.added_kv_proj_dim = added_kv_proj_dim self.added_proj_bias = added_proj_bias self.use_fused_qkv = isinstance(quant_config, NunchakuConfig) self.use_fused_added_qkv = isinstance(quant_config, NunchakuConfig) self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps) if self.use_fused_qkv: self.to_qkv = MergedColumnParallelLinear( query_dim, [self.inner_dim] * 3, bias=bias, gather_output=not self.shard_qkv, 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=not self.shard_qkv, 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=not self.shard_qkv, 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=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.to_v" if prefix else "to_v", ) if not self.pre_only: self.to_out = torch.nn.ModuleList([]) out_proj_cls = RowParallelLinear if self.shard_qkv else ColumnParallelLinear out_proj_kwargs = ( {"input_is_parallel": True} if self.shard_qkv else {"gather_output": True} ) self.to_out.append( out_proj_cls( self.inner_dim, self.out_dim, bias=out_bias, **out_proj_kwargs, quant_config=quant_config, prefix=f"{prefix}.to_out.0" if prefix else "", ) ) if dropout != 0.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: self.to_added_qkv = MergedColumnParallelLinear( added_kv_proj_dim, [self.inner_dim] * 3, bias=added_proj_bias, gather_output=not self.shard_qkv, quant_config=quant_config, 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=not self.shard_qkv, 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=not self.shard_qkv, 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=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.add_v_proj" if prefix else "add_v_proj", ) add_out_proj_cls = ( RowParallelLinear if self.shard_qkv else ColumnParallelLinear ) add_out_proj_kwargs = ( {"input_is_parallel": True} if self.shard_qkv else {"gather_output": True} ) self.to_add_out = add_out_proj_cls( self.inner_dim, query_dim, bias=out_bias, **add_out_proj_kwargs, quant_config=quant_config, prefix=f"{prefix}.to_add_out" if prefix else "", ) self.attn = USPAttention( num_heads=self.local_heads if self.shard_qkv else num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, ) def forward( self, x: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, freqs_cis=None, num_replicated_prefix: int = 0, attn_mask: Optional[torch.Tensor] = None, attn_mask_meta: Optional[Dict[str, int]] = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: ( query, key, value, encoder_query, encoder_key, encoder_value, ) = _get_qkv_projections(self, x, encoder_hidden_states) num_heads = self.local_heads if self.shard_qkv else self.heads query = query.unflatten(-1, (num_heads, -1)) key = key.unflatten(-1, (num_heads, -1)) value = value.unflatten(-1, (num_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, (num_heads, -1)) encoder_key = encoder_key.unflatten(-1, (num_heads, -1)) encoder_value = encoder_value.unflatten(-1, (num_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, ) x = self.attn( query, key, value, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta, num_replicated_prefix=num_replicated_prefix, ) x = x.flatten(2, 3) x = x.to(query.dtype) if encoder_hidden_states is not None: encoder_hidden_states, x = split_seqs( x, encoder_hidden_states.shape[1], sp_txt_pad ) if not self.pre_only: x, _ = self.to_out[0](x) if len(self.to_out) == 2: x = self.to_out[1](x) encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states) return x, encoder_hidden_states else: if not self.pre_only: x, _ = self.to_out[0](x) if len(self.to_out) == 2: x = self.to_out[1](x) return x class FluxSingleTransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) self.use_nunchaku_structure = isinstance(quant_config, NunchakuConfig) self.tp_size = get_tp_world_size() self.local_mlp_hidden_dim = divide(self.mlp_hidden_dim, self.tp_size) self.local_dim = divide(dim, self.tp_size) self.norm = AdaLayerNormZeroSingle(dim) if self.use_nunchaku_structure: self.mlp_fc1 = ColumnParallelLinear( dim, self.mlp_hidden_dim, bias=True, gather_output=True, quant_config=quant_config, prefix=f"{prefix}.mlp_fc1" if prefix else "mlp_fc1", ) self.act_mlp = nn.GELU(approximate="tanh") self.mlp_fc2 = ColumnParallelLinear( self.mlp_hidden_dim, dim, bias=True, gather_output=True, quant_config=quant_config, prefix=f"{prefix}.mlp_fc2" if prefix else "mlp_fc2", ) self.attn = FluxAttention( query_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, bias=True, eps=1e-6, pre_only=False, quant_config=quant_config, prefix=f"{prefix}.attn" if prefix else "attn", ) if is_nunchaku_available(): self.norm = NunchakuAdaLayerNormZeroSingle(self.norm, scale_shift=0) else: shard_single_block = self.tp_size > 1 self.proj_mlp = ColumnParallelLinear( dim, self.mlp_hidden_dim, bias=True, gather_output=not shard_single_block, quant_config=quant_config, prefix=f"{prefix}.proj_mlp" if prefix else "proj_mlp", ) self.act_mlp = nn.GELU(approximate="tanh") proj_out_cls = ( RowParallelLinear if shard_single_block else ColumnParallelLinear ) proj_out_kwargs = ( {"input_is_parallel": True} if shard_single_block else {"gather_output": True} ) self.proj_out = proj_out_cls( dim + self.mlp_hidden_dim, dim, bias=True, **proj_out_kwargs, quant_config=quant_config, prefix=f"{prefix}.proj_out" if prefix else "proj_out", ) if shard_single_block: self._patch_proj_out_weight_loader() self.attn = FluxAttention( query_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, bias=True, eps=1e-6, pre_only=True, quant_config=quant_config, prefix=f"{prefix}.attn" if prefix else "attn", ) def _patch_proj_out_weight_loader(self) -> None: dim, mlp_dim = self.local_dim, self.local_mlp_hidden_dim tp_rank = self.proj_out.tp_rank def _loader(param, loaded_weight): input_dim = getattr(param, "input_dim", None) if input_dim is not None: # checkpoint columns are [attn_full | mlp_full], while TP consumes [attn_shard | mlp_shard] attn_cols = loaded_weight.narrow(input_dim, tp_rank * dim, dim) mlp_cols = loaded_weight.narrow( input_dim, self.tp_size * dim + tp_rank * mlp_dim, mlp_dim, ) param.data.copy_(torch.cat([attn_cols, mlp_cols], dim=input_dim)) else: param.data.copy_(loaded_weight) self.proj_out.weight_loader = _loader if hasattr(self.proj_out.weight, "_weight_loader"): self.proj_out.weight._weight_loader = _loader else: self.proj_out.weight.weight_loader = _loader def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: 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]: text_seq_len = encoder_hidden_states.shape[1] joint_attention_kwargs = joint_attention_kwargs or {} # join_seqs relocates any SP text tail-pad behind the image; the caller # hands single blocks a RoPE cache reordered the same way. sp_txt_pad = (joint_attention_kwargs.get("attn_mask_meta") or {}).get( "local_pad", 0 ) hidden_states = join_seqs(encoder_hidden_states, hidden_states, sp_txt_pad) residual = hidden_states norm_hidden_states, gate = self.norm(hidden_states, emb=temb) if self.use_nunchaku_structure: if _nunchaku_fused_ops_available: mlp_hidden_states = _fused_gelu_mlp( norm_hidden_states, self.mlp_fc1, self.mlp_fc2 ) else: mlp_out, _ = self.mlp_fc1(norm_hidden_states) mlp_hidden_states = self.act_mlp(mlp_out) mlp_hidden_states, _ = self.mlp_fc2(mlp_hidden_states) attn_output = self.attn( x=norm_hidden_states, freqs_cis=freqs_cis, num_replicated_prefix=num_replicated_prefix, **joint_attention_kwargs, ) if isinstance(attn_output, tuple): attn_output = attn_output[0] hidden_states = attn_output + mlp_hidden_states gate = gate.unsqueeze(1) hidden_states = gate * hidden_states hidden_states = residual + hidden_states else: proj_hidden_states, _ = self.proj_mlp(norm_hidden_states) mlp_hidden_states = self.act_mlp(proj_hidden_states) attn_output = self.attn( x=norm_hidden_states, freqs_cis=freqs_cis, num_replicated_prefix=num_replicated_prefix, **joint_attention_kwargs, ) hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) gate = gate.unsqueeze(1) proj_out, _ = self.proj_out(hidden_states) hidden_states = gate * proj_out hidden_states = residual + hidden_states if hidden_states.dtype == torch.float16: hidden_states = hidden_states.clip(-65504, 65504) encoder_hidden_states, hidden_states = split_seqs( hidden_states, text_seq_len, sp_txt_pad ) return encoder_hidden_states, hidden_states class FluxTransformerBlock(nn.Module): def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.norm1 = AdaLayerNormZero(dim) self.norm1_context = AdaLayerNormZero(dim) self.attn = FluxAttention( query_dim=dim, added_kv_proj_dim=dim, dim_head=attention_head_dim, num_heads=num_attention_heads, out_dim=dim, context_pre_only=False, bias=True, eps=eps, quant_config=quant_config, prefix=f"{prefix}.attn" if prefix else "attn", ) self.norm2 = LayerNorm(dim, eps=1e-6, elementwise_affine=False) self.norm2_context = LayerNorm(dim, eps=1e-6, elementwise_affine=False) nunchaku_enabled = ( quant_config is not None and hasattr(quant_config, "get_name") and quant_config.get_name() == "svdquant" and is_nunchaku_available() ) self.use_nunchaku_structure = nunchaku_enabled self.tp_size = get_tp_world_size() if nunchaku_enabled: self.ff = FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate" ) self.ff_context = FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate", ) nunchaku_kwargs = { "precision": quant_config.precision, "rank": quant_config.rank, "act_unsigned": quant_config.act_unsigned, } self.ff = NunchakuFeedForward(self.ff, **nunchaku_kwargs) self.ff_context = NunchakuFeedForward(self.ff_context, **nunchaku_kwargs) self.norm1 = NunchakuAdaLayerNormZero(self.norm1, scale_shift=0) self.norm1_context = NunchakuAdaLayerNormZero( self.norm1_context, scale_shift=0 ) elif self.tp_size > 1: self.ff = FluxParallelFeedForward( dim=dim, dim_out=dim, quant_config=quant_config, prefix=f"{prefix}.ff" if prefix else "ff", ) self.ff_context = FluxParallelFeedForward( dim=dim, dim_out=dim, quant_config=quant_config, prefix=f"{prefix}.ff_context" if prefix else "ff_context", ) else: self.ff = FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate" ) self.ff_context = FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate", ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: 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]: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( hidden_states, emb=temb ) ( norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp, ) = self.norm1_context(encoder_hidden_states, emb=temb) joint_attention_kwargs = joint_attention_kwargs or {} # Attention. attention_outputs = self.attn( x=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, freqs_cis=freqs_cis, num_replicated_prefix=num_replicated_prefix, **joint_attention_kwargs, ) if len(attention_outputs) == 2: attn_output, context_attn_output = attention_outputs elif len(attention_outputs) == 3: attn_output, context_attn_output, ip_attn_output = attention_outputs # Process attention outputs for the `hidden_states`. attn_output = gate_msa.unsqueeze(1) * attn_output hidden_states = hidden_states + attn_output norm_hidden_states = self.norm2(hidden_states) if self.use_nunchaku_structure: norm_hidden_states = ( norm_hidden_states * scale_mlp[:, None] + shift_mlp[:, None] ) else: norm_hidden_states = ( norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] ) ff_output = self.ff(norm_hidden_states) ff_output = gate_mlp.unsqueeze(1) * ff_output hidden_states = hidden_states + ff_output if len(attention_outputs) == 3: hidden_states = hidden_states + ip_attn_output # Process attention outputs for the `encoder_hidden_states`. context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output encoder_hidden_states = encoder_hidden_states + context_attn_output norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) if self.use_nunchaku_structure: norm_encoder_hidden_states = ( norm_encoder_hidden_states * c_scale_mlp[:, None] + c_shift_mlp[:, None] ) else: norm_encoder_hidden_states = ( norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] ) context_ff_output = self.ff_context(norm_encoder_hidden_states) encoder_hidden_states = ( encoder_hidden_states + c_gate_mlp.unsqueeze(1) * 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 FluxPosEmbed(nn.Module): # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 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() 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 FluxTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): """ The Transformer model introduced in Flux. Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ """ param_names_mapping = FluxConfig().arch_config.param_names_mapping @classmethod def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]: return { "skip": [ "norm", "embed", "rotary", "pos_embed", ], "svdq_w4a4": [ "attn.to_qkv", "attn.to_out", "attn.add_qkv_proj", "attn.to_added_qkv", "attn.to_add_out", "img_mlp", "txt_mlp", "attention.to_qkv", "attention.to_out", "proj_mlp", "proj_out", "mlp_fc1", "mlp_fc2", "ff.net", "ff_context.net", ], "awq_w4a16": [ "img_mod", "txt_mod", ], } def __init__( self, config: FluxConfig, hf_config: dict[str, Any], quant_config: Optional[QuantizationConfig] = None, ) -> None: super().__init__(config=config, hf_config=hf_config) self.config = config.arch_config self.out_channels = ( getattr(self.config, "out_channels", None) or self.config.in_channels ) self.inner_dim = ( self.config.num_attention_heads * self.config.attention_head_dim ) self.rotary_emb = FluxPosEmbed(theta=10000, axes_dim=self.config.axes_dims_rope) text_time_guidance_cls = ( CombinedTimestepGuidanceTextProjEmbeddings if self.config.guidance_embeds else CombinedTimestepTextProjEmbeddings ) self.time_text_embed = text_time_guidance_cls( embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim, ) self.context_embedder = ColumnParallelLinear( self.config.joint_attention_dim, self.inner_dim, bias=True, gather_output=True, ) self.x_embedder = ColumnParallelLinear( self.config.in_channels, self.inner_dim, bias=True, gather_output=True ) self.transformer_blocks = nn.ModuleList( [ FluxTransformerBlock( dim=self.inner_dim, num_attention_heads=self.config.num_attention_heads, attention_head_dim=self.config.attention_head_dim, quant_config=quant_config, prefix=f"transformer_blocks.{i}", ) for i in range(self.config.num_layers) ] ) self.single_transformer_blocks = nn.ModuleList( [ FluxSingleTransformerBlock( dim=self.inner_dim, num_attention_heads=self.config.num_attention_heads, attention_head_dim=self.config.attention_head_dim, quant_config=quant_config, prefix=f"single_transformer_blocks.{i}", ) for i in range(self.config.num_single_layers) ] ) self.norm_out = AdaLayerNormContinuous( self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6 ) self.proj_out = ColumnParallelLinear( self.inner_dim, self.config.patch_size * self.config.patch_size * self.out_channels, bias=True, gather_output=True, ) self.layer_names = [ "transformer_blocks", "single_transformer_blocks", ] def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, pooled_projections: torch.Tensor = None, timestep: torch.LongTensor = None, guidance: torch.Tensor = None, freqs_cis: torch.Tensor = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, ) -> Union[torch.Tensor, Transformer2DModelOutput]: """ 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. pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected from the embeddings of input conditions. timestep ( `torch.LongTensor`): Used to indicate denoising step. guidance (`torch.Tensor`): Guidance embeddings. 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). """ if ( joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None ): logger.warning( "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." ) hidden_states, _ = self.x_embedder(hidden_states) # Only pass guidance to time_text_embed if the model supports it if self.config.guidance_embeds and guidance is not None: temb = self.time_text_embed(timestep, guidance, pooled_projections) else: temb = self.time_text_embed(timestep, pooled_projections) num_txt_tokens = encoder_hidden_states.shape[1] 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 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 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 # Single blocks apply RoPE on the relocated [txt_real, img, pad] # layout, so hand them a cache reordered the same way. if freqs_cis is not None: t_loc = txt_shard.local_len pad = txt_shard.local_pad singles_freqs_cis = ( join_seqs(cos[:t_loc], cos[t_loc:], pad, dim=0), join_seqs(sin[:t_loc], sin[t_loc:], pad, dim=0), ) if ( joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs ): ip_adapter_image_embeds = joint_attention_kwargs.pop( "ip_adapter_image_embeds" ) ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds) joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states}) for block in self.transformer_blocks: encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, freqs_cis=freqs_cis, joint_attention_kwargs=joint_attention_kwargs, num_replicated_prefix=num_replicated_prefix, ) for block in self.single_transformer_blocks: encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, freqs_cis=singles_freqs_cis, joint_attention_kwargs=joint_attention_kwargs, num_replicated_prefix=num_replicated_prefix, ) hidden_states = self.norm_out(hidden_states, temb) output, _ = self.proj_out(hidden_states) return output EntryClass = FluxTransformer2DModel