flux_2

class TimestepEmbedding(nn.Module):
def __init__(
    self,
    in_channels: int,
    time_embed_dim: int,
    out_dim: Optional[int] = None,
    post_act_fn: Optional[str] = None,
    cond_proj_dim: Optional[int] = None,
    sample_proj_bias: bool = True,
):
    super().__init__()
    self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
    self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) if cond_proj_dim is not None else None
    self.act = nn.SiLU()
    time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim
    self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
    self.post_act = None if post_act_fn is None else get_activation(post_act_fn)

def forward(self, sample: torch.Tensor, condition: Optional[torch.Tensor] = None) -> torch.Tensor:
    if condition is not None and self.cond_proj is not None:
        sample = sample + self.cond_proj(condition)
    sample = self.linear_1(sample)
    if self.act is not None:
        sample = self.act(sample)
    sample = self.linear_2(sample)
    if self.post_act is not None:
        sample = self.post_act(sample)
    return sample</code></pre></details>
fastvideo.models.dits.flux_2.TimestepEmbedding.forward
forward(sample: torch.Tensor, condition: Optional[torch.Tensor] = None) -> torch.Tensor
Source
def forward(self, sample: torch.Tensor, condition: Optional[torch.Tensor] = None) -> torch.Tensor:
    if condition is not None and self.cond_proj is not None:
        sample = sample + self.cond_proj(condition)
    sample = self.linear_1(sample)
    if self.act is not None:
        sample = self.act(sample)
    sample = self.linear_2(sample)
    if self.post_act is not None:
        sample = self.post_act(sample)
    return sample
fastvideo.models.dits.flux_2.AdaLayerNormContinuous
AdaLayerNormContinuous(embedding_dim: int, conditioning_embedding_dim: int, elementwise_affine: bool = True, eps: float = 1e-05, bias: bool = True, norm_type: str = 'layer_norm')
Source
class AdaLayerNormContinuous(nn.Module):
def __init__(
    self,
    embedding_dim: int,
    conditioning_embedding_dim: int,
    elementwise_affine: bool = True,
    eps: float = 1e-5,
    bias: bool = True,
    norm_type: str = &quot;layer_norm&quot;,
):
    super().__init__()
    self.silu = nn.SiLU()
    self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
    if norm_type != &quot;layer_norm&quot;:
        raise ValueError(f&quot;unknown norm_type &#123;norm_type&#125;&quot;)
    self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine, bias=bias)

def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -&gt; torch.Tensor:
    emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
    scale, shift = torch.chunk(emb, 2, dim=1)
    return self.norm(x) * (1 + scale)&#91;:, None, :&#93; + shift&#91;:, None, :&#93;</code></pre></details>
fastvideo.models.dits.flux_2.AdaLayerNormContinuous.forward
forward(x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
    emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
    scale, shift = torch.chunk(emb, 2, dim=1)
    return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
fastvideo.models.dits.flux_2.apply_qk_norm
apply_qk_norm(q: torch.Tensor, k: torch.Tensor, q_norm: nn.Module, k_norm: nn.Module, head_dim: int, allow_inplace: bool = True) -> Tuple[torch.Tensor, torch.Tensor]
Apply QK normalization for query and key tensors.
Source
def apply_qk_norm(
    q: torch.Tensor,
    k: torch.Tensor,
    q_norm: nn.Module,
    k_norm: nn.Module,
    head_dim: int,
    allow_inplace: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Apply QK normalization for query and key tensors."""
    q_shape = q.shape
    k_shape = k.shape
    q_out = q_norm(q.view(-1, head_dim)).view(q_shape)
    k_out = k_norm(k.view(-1, head_dim)).view(k_shape)
    return q_out, k_out
fastvideo.models.dits.flux_2.Flux2SwiGLU
Flux2SwiGLU(compute_in_fp32: bool = False)
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.
Source
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, compute_in_fp32: bool = False):
    super().__init__()
    self.gate_fn = nn.SiLU()
    self.compute_in_fp32 = compute_in_fp32

def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    if self.compute_in_fp32:
        return (self.gate_fn(x1.float()) * x2.float()).to(x.dtype)
    return self.gate_fn(x1) * x2</code></pre></details>
fastvideo.models.dits.flux_2.Flux2SwiGLU.forward
forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    if self.compute_in_fp32:
        return (self.gate_fn(x1.float()) * x2.float()).to(x.dtype)
    return self.gate_fn(x1) * x2
fastvideo.models.dits.flux_2.Flux2FeedForward
Flux2FeedForward(dim: int, dim_out: Optional[int] = None, mult: float = 3.0, inner_dim: Optional[int] = None, bias: bool = False, swiglu_fp32: bool = False)
FF sub-block: SwiGLU + two linears.
At tp_world_size == 1, use nn.Linear to match Diffusers numerics and
checkpoint layout; under tensor parallel, use ColumnParallelLinear.
Source
class Flux2FeedForward(nn.Module):
    """FF sub-block: SwiGLU + two linears.
At ``tp_world_size == 1``, use ``nn.Linear`` to match Diffusers numerics and
checkpoint layout; under tensor parallel, use ``ColumnParallelLinear``.
&quot;&quot;&quot;

def __init__(
    self,
    dim: int,
    dim_out: Optional&#91;int&#93; = None,
    mult: float = 3.0,
    inner_dim: Optional&#91;int&#93; = None,
    bias: bool = False,
    swiglu_fp32: bool = False,
):
    super().__init__()
    if inner_dim is None:
        inner_dim = int(dim * mult)
    dim_out = dim_out or dim

    self._ff_dense_linear = get_tp_world_size() == 1
    if self._ff_dense_linear:
        self.linear_in = nn.Linear(dim, inner_dim * 2, bias=bias)
        self.linear_out = nn.Linear(inner_dim, dim_out, bias=bias)
    else:
        self.linear_in = ColumnParallelLinear(dim, inner_dim * 2, bias=bias, gather_output=True)
        self.linear_out = ColumnParallelLinear(inner_dim, dim_out, bias=bias, gather_output=True)
    self.act_fn = Flux2SwiGLU(compute_in_fp32=swiglu_fp32)

def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
    if self._ff_dense_linear:
        x = self.linear_in(x)
        x = self.act_fn(x)
        x = self.linear_out(x)
    else:
        x, _ = self.linear_in(x)
        x = self.act_fn(x)
        x, _ = self.linear_out(x)
    return x</code></pre></details>
fastvideo.models.dits.flux_2.Flux2FeedForward.forward
forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    if self._ff_dense_linear:
        x = self.linear_in(x)
        x = self.act_fn(x)
        x = self.linear_out(x)
    else:
        x, _ = self.linear_in(x)
        x = self.act_fn(x)
        x, _ = self.linear_out(x)
    return x
fastvideo.models.dits.flux_2.Flux2Attention
Flux2Attention(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-06, out_dim: int = None, elementwise_affine: bool = True, supported_attention_backends: Optional[Tuple[AttentionBackendEnum, ...]] = None)
Joint attention for Flux2 double-stream blocks (image + text).
Context (text) branch uses add_q_proj / add_k_proj / add_v_proj on
encoder_hidden_states, QK norm via norm_added_*, concatenation text then
image on sequence dim, RoPE on the full sequence, attention, then
to_add_out on the text slice only. Diffusers does the same layout in
Flux2AttnProcessor with nn.Linear; here those linears are
ColumnParallelLinear (F.linear at tp_size==1). For parity debugging,
see compare_flux2_context_branch_double0.py.
Source
class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
    """Joint attention for Flux2 double-stream blocks (image + text).
**Context (text) branch** uses ``add_q_proj`` / ``add_k_proj`` / ``add_v_proj`` on
``encoder_hidden_states``, QK norm via ``norm_added_*``, concatenation **text then
image** on sequence dim, RoPE on the full sequence, attention, then
``to_add_out`` on the **text** slice only. Diffusers does the same layout in
``Flux2AttnProcessor`` with ``nn.Linear``; here those linears are
``ColumnParallelLinear`` (``F.linear`` at ``tp_size==1``). For parity debugging,
see ``compare_flux2_context_branch_double0.py``.
&quot;&quot;&quot;

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&#91;int&#93; = None,
    added_proj_bias: Optional&#91;bool&#93; = True,
    out_bias: bool = True,
    eps: float = 1e-6,
    out_dim: int = None,
    elementwise_affine: bool = True,
    supported_attention_backends: Optional&#91;Tuple&#91;AttentionBackendEnum, ...&#93;&#93; = None,
):
    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.use_bias = bias
    self.dropout = dropout

    self.added_kv_proj_dim = added_kv_proj_dim
    self.added_proj_bias = added_proj_bias

    self.to_q = ColumnParallelLinear(query_dim, self.inner_dim, bias=bias, gather_output=True)
    self.to_k = ColumnParallelLinear(query_dim, self.inner_dim, bias=bias, gather_output=True)
    self.to_v = ColumnParallelLinear(query_dim, self.inner_dim, bias=bias, gather_output=True)

    # QK norm: match Diffusers ``nn.RMSNorm`` (not vLLM-style fp32 variance path).
    self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
    self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)

    self.to_out = torch.nn.ModuleList(&#91;&#93;)
    self.to_out.append(ColumnParallelLinear(self.inner_dim, self.out_dim, bias=out_bias, gather_output=True))
    self.to_out.append(torch.nn.Dropout(dropout))

    if added_kv_proj_dim is not None:
        self.norm_added_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
        self.norm_added_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
        self.add_q_proj = ColumnParallelLinear(
            added_kv_proj_dim,
            self.inner_dim,
            bias=added_proj_bias,
            gather_output=True,
        )
        self.add_k_proj = ColumnParallelLinear(
            added_kv_proj_dim,
            self.inner_dim,
            bias=added_proj_bias,
            gather_output=True,
        )
        self.add_v_proj = ColumnParallelLinear(
            added_kv_proj_dim,
            self.inner_dim,
            bias=added_proj_bias,
            gather_output=True,
        )
        self.to_add_out = ColumnParallelLinear(self.inner_dim, query_dim, bias=out_bias, gather_output=True)

    self.attn = LocalAttention(
        num_heads=num_heads,
        head_size=self.head_dim,
        softmax_scale=None,
        causal=False,
        supported_attention_backends=supported_attention_backends,
    )

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: Optional&#91;torch.Tensor&#93; = None,
    freqs_cis: Optional&#91;Tuple&#91;torch.Tensor, torch.Tensor&#93;&#93; = None,
) -&gt; 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.heads, -1))
    key = key.unflatten(-1, (self.heads, -1))
    value = value.unflatten(-1, (self.heads, -1))

    query, key = apply_qk_norm(
        q=query,
        k=key,
        q_norm=self.norm_q,
        k_norm=self.norm_k,
        head_dim=self.head_dim,
        allow_inplace=True,
    )

    if self.added_kv_proj_dim is not None:
        encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
        encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
        encoder_value = encoder_value.unflatten(-1, (self.heads, -1))

        encoder_query, encoder_key = apply_qk_norm(
            q=encoder_query,
            k=encoder_key,
            q_norm=self.norm_added_q,
            k_norm=self.norm_added_k,
            head_dim=self.head_dim,
            allow_inplace=True,
        )

        query = torch.cat(&#91;encoder_query, query&#93;, dim=1)
        key = torch.cat(&#91;encoder_key, key&#93;, dim=1)
        value = torch.cat(&#91;encoder_value, value&#93;, dim=1)

    if freqs_cis is not None:
        cos, sin = freqs_cis
        # Match Diffusers: keep cos/sin in fp32 for RoPE math (do not cast to bf16).
        cos = cos.to(device=query.device)
        sin = sin.to(device=query.device)
        # Diffusers Flux2: &#91;B, S, H, D&#93; with sequence_dim=1 (cos/sin &#91;S, D&#93;).
        query = apply_rotary_emb(
            query,
            (cos, sin),
            use_real=True,
            use_real_unbind_dim=-1,
            sequence_dim=1,
        )
        key = apply_rotary_emb(
            key,
            (cos, sin),
            use_real=True,
            use_real_unbind_dim=-1,
            sequence_dim=1,
        )

    hidden_states = self.attn(query, key, value)

    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 = hidden_states.split_with_sizes(
            &#91;
                encoder_hidden_states.shape&#91;1&#93;,
                hidden_states.shape&#91;1&#93; - encoder_hidden_states.shape&#91;1&#93;,
            &#93;,
            dim=1,
        )
        encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)

    hidden_states, _ = self.to_out&#91;0&#93;(hidden_states)
    hidden_states = self.to_out&#91;1&#93;(hidden_states)

    if encoder_hidden_states is not None:
        return hidden_states, encoder_hidden_states
    else:
        return hidden_states</code></pre></details>
fastvideo.models.dits.flux_2.Flux2Attention.forward
forward(hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: Optional[torch.Tensor] = None,
    freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = 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.heads, -1))
key = key.unflatten(-1, (self.heads, -1))
value = value.unflatten(-1, (self.heads, -1))

query, key = apply_qk_norm(
    q=query,
    k=key,
    q_norm=self.norm_q,
    k_norm=self.norm_k,
    head_dim=self.head_dim,
    allow_inplace=True,
)

if self.added_kv_proj_dim is not None:
    encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
    encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
    encoder_value = encoder_value.unflatten(-1, (self.heads, -1))

    encoder_query, encoder_key = apply_qk_norm(
        q=encoder_query,
        k=encoder_key,
        q_norm=self.norm_added_q,
        k_norm=self.norm_added_k,
        head_dim=self.head_dim,
        allow_inplace=True,
    )

    query = torch.cat(&#91;encoder_query, query&#93;, dim=1)
    key = torch.cat(&#91;encoder_key, key&#93;, dim=1)
    value = torch.cat(&#91;encoder_value, value&#93;, dim=1)

if freqs_cis is not None:
    cos, sin = freqs_cis
    # Match Diffusers: keep cos/sin in fp32 for RoPE math (do not cast to bf16).
    cos = cos.to(device=query.device)
    sin = sin.to(device=query.device)
    # Diffusers Flux2: &#91;B, S, H, D&#93; with sequence_dim=1 (cos/sin &#91;S, D&#93;).
    query = apply_rotary_emb(
        query,
        (cos, sin),
        use_real=True,
        use_real_unbind_dim=-1,
        sequence_dim=1,
    )
    key = apply_rotary_emb(
        key,
        (cos, sin),
        use_real=True,
        use_real_unbind_dim=-1,
        sequence_dim=1,
    )

hidden_states = self.attn(query, key, value)

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 = hidden_states.split_with_sizes(
        &#91;
            encoder_hidden_states.shape&#91;1&#93;,
            hidden_states.shape&#91;1&#93; - encoder_hidden_states.shape&#91;1&#93;,
        &#93;,
        dim=1,
    )
    encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)

hidden_states, _ = self.to_out&#91;0&#93;(hidden_states)
hidden_states = self.to_out&#91;1&#93;(hidden_states)

if encoder_hidden_states is not None:
    return hidden_states, encoder_hidden_states
else:
    return hidden_states</code></pre></details>
fastvideo.models.dits.flux_2.Flux2ParallelSelfAttention
Flux2ParallelSelfAttention(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-06, out_dim: int = None, elementwise_affine: bool = True, mlp_ratio: float = 4.0, mlp_mult_factor: int = 2, supported_attention_backends: Optional[Tuple[AttentionBackendEnum, ...]] = None)
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 for a visual depiction of this type of transformer block.
Source
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 &#91;ViT-22B
paper&#93;(https://arxiv.org/abs/2302.05442) for a visual depiction of this type of transformer block.
&quot;&quot;&quot;

# 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-6,
    out_dim: int = None,
    elementwise_affine: bool = True,
    mlp_ratio: float = 4.0,
    mlp_mult_factor: int = 2,
    supported_attention_backends: Optional&#91;Tuple&#91;AttentionBackendEnum, ...&#93;&#93; = None,
):
    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.use_bias = bias
    self.dropout = dropout

    self.mlp_ratio = mlp_ratio
    self.mlp_hidden_dim = int(query_dim * self.mlp_ratio)
    self.mlp_mult_factor = mlp_mult_factor

    # Fused QKV projections + MLP input projection
    self.to_qkv_mlp_proj = ColumnParallelLinear(
        self.query_dim,
        self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
        bias=bias,
        gather_output=True,
    )
    self.mlp_act_fn = Flux2SwiGLU(compute_in_fp32=False)

    self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
    self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)

    # Fused attention output projection + MLP output projection
    self.to_out = ColumnParallelLinear(
        self.inner_dim + self.mlp_hidden_dim,
        self.out_dim,
        bias=out_bias,
        gather_output=True,
    )

    self.attn = LocalAttention(
        num_heads=num_heads,
        head_size=self.head_dim,
        softmax_scale=None,
        causal=False,
        supported_attention_backends=supported_attention_backends,
    )

def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional&#91;torch.Tensor&#93; = None,
    freqs_cis: Optional&#91;Tuple&#91;torch.Tensor, torch.Tensor&#93;&#93; = None,
    **kwargs,
) -&gt; torch.Tensor:
    # Parallel in (QKV + MLP in) projection
    hidden_states, _ = self.to_qkv_mlp_proj(hidden_states)
    qkv, mlp_hidden_states = torch.split(
        hidden_states,
        &#91;3 * self.inner_dim, self.mlp_hidden_dim * self.mlp_mult_factor&#93;,
        dim=-1,
    )

    # Handle the attention logic
    query, key, value = qkv.chunk(3, dim=-1)

    query = query.unflatten(-1, (self.heads, -1))
    key = key.unflatten(-1, (self.heads, -1))
    value = value.unflatten(-1, (self.heads, -1))

    query = self.norm_q(query)
    key = self.norm_k(key)

    if freqs_cis is not None:
        cos, sin = freqs_cis
        cos = cos.to(device=query.device)
        sin = sin.to(device=query.device)
        # Single stream: match diffusers sequence_dim=1 (query/key &#91;B, S, H, D&#93;, no transpose)
        query = apply_rotary_emb(query, (cos, sin), use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
        key = apply_rotary_emb(key, (cos, sin), use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
    hidden_states = self.attn(query, key, value)
    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(&#91;hidden_states, mlp_hidden_states&#93;, dim=-1)
    hidden_states, _ = self.to_out(hidden_states)

    return hidden_states</code></pre></details>
fastvideo.models.dits.flux_2.Flux2ParallelSelfAttention.forward
forward(hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    **kwargs,
) -> torch.Tensor:
    # 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.inner_dim, self.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.heads, -1))
key = key.unflatten(-1, (self.heads, -1))
value = value.unflatten(-1, (self.heads, -1))

query = self.norm_q(query)
key = self.norm_k(key)

if freqs_cis is not None:
    cos, sin = freqs_cis
    cos = cos.to(device=query.device)
    sin = sin.to(device=query.device)
    # Single stream: match diffusers sequence_dim=1 (query/key &#91;B, S, H, D&#93;, no transpose)
    query = apply_rotary_emb(query, (cos, sin), use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
    key = apply_rotary_emb(key, (cos, sin), use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
hidden_states = self.attn(query, key, value)
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(&#91;hidden_states, mlp_hidden_states&#93;, dim=-1)
hidden_states, _ = self.to_out(hidden_states)

return hidden_states</code></pre></details>
fastvideo.models.dits.flux_2.Flux2SingleTransformerBlock
Flux2SingleTransformerBlock(dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 3.0, eps: float = 1e-06, bias: bool = False, supported_attention_backends: Optional[Tuple[AttentionBackendEnum, ...]] = None)
Source
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: Optional&#91;Tuple&#91;AttentionBackendEnum, ...&#93;&#93; = None,
):
    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 &quot;parallel&quot; transformer block. See the &#91;ViT-22B paper&#93;(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,
    )

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: Optional&#91;torch.Tensor&#93;,
    temb_mod_params: Tuple&#91;torch.Tensor, torch.Tensor, torch.Tensor&#93;,
    freqs_cis: Optional&#91;Tuple&#91;torch.Tensor, torch.Tensor&#93;&#93; = None,
    joint_attention_kwargs: Optional&#91;Dict&#91;str, Any&#93;&#93; = None,
    split_hidden_states: bool = False,
    text_seq_len: Optional&#91;int&#93; = None,
) -&gt; Tuple&#91;torch.Tensor, torch.Tensor&#93;:
    # 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&#91;1&#93;
        hidden_states = torch.cat(&#91;encoder_hidden_states, hidden_states&#93;, 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 &#123;&#125;
    attn_output = self.attn(
        hidden_states=norm_hidden_states,
        freqs_cis=freqs_cis,
        **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&#91;:, :text_seq_len&#93;,
            hidden_states&#91;:, text_seq_len:&#93;,
        )
        return encoder_hidden_states, hidden_states
    else:
        return hidden_states</code></pre></details>
fastvideo.models.dits.flux_2.Flux2SingleTransformerBlock.forward
forward(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) -> Tuple[torch.Tensor, torch.Tensor]
Source
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,
) -> 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 &#123;&#125;
attn_output = self.attn(
    hidden_states=norm_hidden_states,
    freqs_cis=freqs_cis,
    **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&#91;:, :text_seq_len&#93;,
        hidden_states&#91;:, text_seq_len:&#93;,
    )
    return encoder_hidden_states, hidden_states
else:
    return hidden_states</code></pre></details>
fastvideo.models.dits.flux_2.Flux2TransformerBlock
Flux2TransformerBlock(dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 3.0, eps: float = 1e-06, bias: bool = False, supported_attention_backends: Optional[Tuple[AttentionBackendEnum, ...]] = None, ff_context_swiglu_fp32: bool = False)
Source
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: Optional&#91;Tuple&#91;AttentionBackendEnum, ...&#93;&#93; = None,
    ff_context_swiglu_fp32: bool = False,
):
    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,
    )

    self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
    self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias)

    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,
        swiglu_fp32=ff_context_swiglu_fp32,
    )

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    temb_mod_params_img: Tuple&#91;Tuple&#91;torch.Tensor, torch.Tensor, torch.Tensor&#93;, ...&#93;,
    temb_mod_params_txt: Tuple&#91;Tuple&#91;torch.Tensor, torch.Tensor, torch.Tensor&#93;, ...&#93;,
    freqs_cis: Optional&#91;Tuple&#91;torch.Tensor, torch.Tensor&#93;&#93; = None,
    joint_attention_kwargs: Optional&#91;Dict&#91;str, Any&#93;&#93; = None,
) -&gt; Tuple&#91;torch.Tensor, torch.Tensor&#93;:
    joint_attention_kwargs = joint_attention_kwargs or &#123;&#125;

    # Modulation parameters shape: &#91;1, 1, self.dim&#93;
    (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,
    )

    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)
    context_ff_update = c_gate_mlp * context_ff_output
    encoder_hidden_states = encoder_hidden_states + context_ff_update
    if encoder_hidden_states.dtype == torch.float16:
        encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)

    return encoder_hidden_states, hidden_states</code></pre></details>
fastvideo.models.dits.flux_2.Flux2TransformerBlock.forward
forward(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) -> Tuple[torch.Tensor, torch.Tensor]
Source
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,
) -> Tuple[torch.Tensor, torch.Tensor]:
    joint_attention_kwargs = joint_attention_kwargs or {}
# Modulation parameters shape: &#91;1, 1, self.dim&#93;
(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,
)

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)
context_ff_update = c_gate_mlp * context_ff_output
encoder_hidden_states = encoder_hidden_states + context_ff_update
if encoder_hidden_states.dtype == torch.float16:
    encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)

return encoder_hidden_states, hidden_states</code></pre></details>
fastvideo.models.dits.flux_2.Flux2TimestepGuidanceEmbeddings
Flux2TimestepGuidanceEmbeddings(in_channels: int = 256, embedding_dim: int = 6144, bias: bool = False, guidance_embeds: bool = True)
Source
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&#91;torch.Tensor&#93; = None) -&gt; 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</code></pre></details>
fastvideo.models.dits.flux_2.Flux2TimestepGuidanceEmbeddings.forward
forward(timestep: torch.Tensor, guidance: Optional[torch.Tensor] = None) -> torch.Tensor
Source
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</code></pre></details>
fastvideo.models.dits.flux_2.Flux2Modulation
Flux2Modulation(dim: int, mod_param_sets: int = 2, bias: bool = False)
Source
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) -&gt; Tuple&#91;Tuple&#91;torch.Tensor, torch.Tensor, torch.Tensor&#93;, ...&#93;:
    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&#91;3 * i:3 * (i + 1)&#93; for i in range(self.mod_param_sets))</code></pre></details>
fastvideo.models.dits.flux_2.Flux2Modulation.forward
forward(temb: torch.Tensor) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]
Source
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&#91;3 * i:3 * (i + 1)&#93; for i in range(self.mod_param_sets))</code></pre></details>
fastvideo.models.dits.flux_2.Flux2PosEmbed
Flux2PosEmbed(theta: int, axes_dim: List[int])
Flux2 N-D RoPE: matches Diffusers Flux2PosEmbed (transformer_flux2).
Source
class Flux2PosEmbed(nn.Module):
    """Flux2 N-D RoPE: matches Diffusers ``Flux2PosEmbed`` (transformer_flux2)."""
def __init__(self, theta: int, axes_dim: List&#91;int&#93;):
    super().__init__()
    self.theta = theta
    self.axes_dim = axes_dim

def forward_uncached(self, pos: torch.Tensor) -&gt; tuple&#91;torch.Tensor, torch.Tensor&#93;:
    &quot;&quot;&quot;Same as ``forward``; kept for callers that pre-flatten positions.&quot;&quot;&quot;
    return self.forward(pos)

def forward(self, ids: torch.Tensor) -&gt; tuple&#91;torch.Tensor, torch.Tensor&#93;:
    &quot;&quot;&quot;Per Diffusers: loop axes, 1D RoPE each, concat on channel dim.&quot;&quot;&quot;
    pos = ids.float()
    is_mps = pos.device.type == &quot;mps&quot;
    is_npu = pos.device.type == &quot;npu&quot;
    freqs_dtype = (torch.float32 if (is_mps or is_npu) else torch.float64)
    cos_out: List&#91;torch.Tensor&#93; = &#91;&#93;
    sin_out: List&#91;torch.Tensor&#93; = &#91;&#93;
    for i in range(len(self.axes_dim)):
        c, s = _flux2_get_1d_rotary_pos_embed_diffusers(
            self.axes_dim&#91;i&#93;,
            pos&#91;..., i&#93;,
            float(self.theta),
            freqs_dtype,
        )
        cos_out.append(c)
        sin_out.append(s)
    freqs_cos = torch.cat(cos_out, dim=-1).to(device=ids.device)
    freqs_sin = torch.cat(sin_out, dim=-1).to(device=ids.device)
    return freqs_cos.contiguous().float(), freqs_sin.contiguous().float()</code></pre></details>
fastvideo.models.dits.flux_2.Flux2PosEmbed.forward_uncached
forward_uncached(pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]
Same as forward; kept for callers that pre-flatten positions.
Source
def forward_uncached(self, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Same as ``forward``; kept for callers that pre-flatten positions."""
    return self.forward(pos)
fastvideo.models.dits.flux_2.Flux2PosEmbed.forward
forward(ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]
Per Diffusers: loop axes, 1D RoPE each, concat on channel dim.
Source
def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Per Diffusers: loop axes, 1D RoPE each, concat on channel dim."""
    pos = ids.float()
    is_mps = pos.device.type == "mps"
    is_npu = pos.device.type == "npu"
    freqs_dtype = (torch.float32 if (is_mps or is_npu) else torch.float64)
    cos_out: List[torch.Tensor] = []
    sin_out: List[torch.Tensor] = []
    for i in range(len(self.axes_dim)):
        c, s = _flux2_get_1d_rotary_pos_embed_diffusers(
            self.axes_dim[i],
            pos[..., i],
            float(self.theta),
            freqs_dtype,
        )
        cos_out.append(c)
        sin_out.append(s)
    freqs_cos = torch.cat(cos_out, dim=-1).to(device=ids.device)
    freqs_sin = torch.cat(sin_out, dim=-1).to(device=ids.device)
    return freqs_cos.contiguous().float(), freqs_sin.contiguous().float()
fastvideo.models.dits.flux_2.compute_flux2_freqs_cis_from_ids
compute_flux2_freqs_cis_from_ids(rotary_emb: Flux2PosEmbed, text_ids: torch.Tensor, latent_ids: torch.Tensor, device: torch.device, dtype: torch.dtype | None = None) -> tuple[torch.Tensor, torch.Tensor]
Build (cos, sin) like Diffusers: pos_embed(txt) then pos_embed(img), cat dim 0.
Returns cos/sin in fp32 on device (dtype is ignored) so RoPE matches Diffusers.
Source
def compute_flux2_freqs_cis_from_ids(
    rotary_emb: Flux2PosEmbed,
    text_ids: torch.Tensor,
    latent_ids: torch.Tensor,
    device: torch.device,
    dtype: torch.dtype | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Build (cos, sin) like Diffusers: ``pos_embed(txt)`` then ``pos_embed(img)``, cat dim 0.
Returns cos/sin in **fp32** on ``device`` (``dtype`` is ignored) so RoPE matches Diffusers.
&quot;&quot;&quot;
tid = text_ids
lid = latent_ids
if tid.ndim == 3:
    tid = tid&#91;0&#93;
if lid.ndim == 3:
    lid = lid&#91;0&#93;
n_axes = len(rotary_emb.axes_dim)
tid = _flux2_pad_ids_to_axes(tid, n_axes)
lid = _flux2_pad_ids_to_axes(lid, n_axes)
with torch.no_grad():
    text_cos, text_sin = rotary_emb.forward(tid)
    img_cos, img_sin = rotary_emb.forward(lid)
cos = torch.cat(&#91;text_cos, img_cos&#93;, dim=0).to(device=device)
sin = torch.cat(&#91;text_sin, img_sin&#93;, dim=0).to(device=device)
# Match Diffusers: keep RoPE cos/sin in fp32 (``dtype`` unused; callers may still pass it).
_ = dtype
return (cos, sin)</code></pre></details>
fastvideo.models.dits.flux_2.Flux2Transformer2DModel
Flux2Transformer2DModel(config: Flux2Config, hf_config: dict[str, Any])
The Transformer model introduced in Flux 2.
Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
Source
class Flux2Transformer2DModel(BaseDiT):
    """
    The Transformer model introduced in Flux 2.
Reference: https://blackforestlabs.ai/announcing-black-forest-labs/

&quot;&quot;&quot;

param_names_mapping = Flux2Config().param_names_mapping
_fsdp_shard_conditions = Flux2Config()._fsdp_shard_conditions
_compile_conditions = Flux2Config()._compile_conditions

def __init__(self, config: Flux2Config, hf_config: dict&#91;str, Any&#93;):
    super().__init__(config=config, hf_config=hf_config)
    patch_size: int = config.patch_size
    in_channels: int = config.in_channels
    out_channels: Optional&#91;int&#93; = 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&#91;int, ...&#93; = config.axes_dims_rope
    rope_theta: int = config.rope_theta
    eps: float = config.eps
    guidance_embeds: bool = getattr(config, &quot;guidance_embeds&quot;, True)
    self.out_channels = out_channels or in_channels
    self.inner_dim = num_attention_heads * attention_head_dim
    self.guidance_embeds = guidance_embeds

    # Add required BaseDiT attributes
    self.hidden_size = self.inner_dim
    self.num_attention_heads = num_attention_heads
    self.num_channels_latents = self.out_channels

    # 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
    supported_attention_backends = self._supported_attention_backends
    ff_ctx_swiglu = getattr(config.arch_config, &quot;ff_context_swiglu_fp32&quot;, False)
    self.transformer_blocks = nn.ModuleList(&#91;
        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=supported_attention_backends,
            ff_context_swiglu_fp32=ff_ctx_swiglu,
        ) for _ in range(num_layers)
    &#93;)

    # 6. Single Stream Transformer Blocks
    self.single_transformer_blocks = nn.ModuleList(&#91;
        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=supported_attention_backends,
        ) for _ in range(num_single_layers)
    &#93;)

    # 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,
    )

    self.layer_names = &#91;&quot;transformer_blocks&quot;, &quot;single_transformer_blocks&quot;&#93;

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor = None,
    timestep: torch.LongTensor = None,
    img_ids: torch.Tensor = None,
    txt_ids: torch.Tensor = None,
    guidance: torch.Tensor = None,
    freqs_cis: torch.Tensor = None,
    joint_attention_kwargs: Optional&#91;Dict&#91;str, Any&#93;&#93; = None,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;
    The &#91;`Flux2Transformer2DModel`&#93; 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
            &#91;diffusers.models.attention_processor&#93;(https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).

    &quot;&quot;&quot;
    # 0. Handle input arguments
    joint_attention_kwargs = joint_attention_kwargs.copy() if joint_attention_kwargs is not None else &#123;&#125;

    # Pipeline passes prompt_embeds as a list (one per text encoder); use first
    if isinstance(encoder_hidden_states, (list, tuple)):
        if len(encoder_hidden_states) &gt; 1:
            logger.warning(
                &quot;Flux 2 received %d encoder outputs but uses only the first; &quot;
                &quot;indices &gt;= 1 are dropped. Wire multi-encoder support or pass only one.&quot;,
                len(encoder_hidden_states),
            )
        encoder_hidden_states = encoder_hidden_states&#91;0&#93;

    num_txt_tokens = encoder_hidden_states.shape&#91;1&#93;

    # 0.5. Reshape 5D latent (B, C, T, H, W) from pipeline to 3D (B, T*H*W, C) for transformer
    input_was_5d = False
    if hidden_states.dim() == 5:
        input_was_5d = True
        b, c_in, t, h, w = hidden_states.shape
        hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape(b, t * h * w, c_in)

    # 1. Calculate timestep embedding and modulation parameters (match diffusers: scale to &#91;0, 1000&#93;)
    timestep = timestep.to(hidden_states.dtype) * 1000
    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)&#91;0&#93;

    # 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)

    # 3. Compute RoPE positional embeddings when not provided externally
    if freqs_cis is None:
        if txt_ids is None or img_ids is None:
            if input_was_5d:
                img_h, img_w = h, w
            else:
                raise ValueError(&quot;Flux 2 4D input requires explicit txt_ids and img_ids &quot;
                                 &quot;with real image dimensions; square-image fallback removed.&quot;)

            txt_len = encoder_hidden_states.shape&#91;1&#93;
            txt_ids = torch.cartesian_prod(
                torch.arange(1),
                torch.arange(1),
                torch.arange(1),
                torch.arange(txt_len),
            ).to(device=hidden_states.device)
            img_ids = torch.cartesian_prod(
                torch.arange(1),
                torch.arange(img_h),
                torch.arange(img_w),
                torch.arange(1),
            ).to(device=hidden_states.device)
        freqs_cis = compute_flux2_freqs_cis_from_ids(
            self.rotary_emb,
            txt_ids,
            img_ids,
            device=hidden_states.device,
        )

    # 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,
        )
    # Concatenate text and image streams for single-block inference
    hidden_states = torch.cat(&#91;encoder_hidden_states, hidden_states&#93;, dim=1)

    # 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=freqs_cis,
            joint_attention_kwargs=joint_attention_kwargs,
        )
    # Remove text tokens from concatenated stream
    hidden_states = hidden_states&#91;:, num_txt_tokens:, ...&#93;

    # 6. Output layers
    hidden_states = self.norm_out(hidden_states, temb)
    output, _ = self.proj_out(hidden_states)

    # Reshape 3D (B, T*H*W, C) back to 5D (B, C, T, H, W) for scheduler
    if input_was_5d:
        output = output.reshape(b, t, h, w, self.out_channels).permute(0, 4, 1, 2, 3)

    return output</code></pre></details>
fastvideo.models.dits.flux_2.Flux2Transformer2DModel.forward
forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor = None, timestep: torch.LongTensor = None, img_ids: torch.Tensor = None, txt_ids: torch.Tensor = None, guidance: torch.Tensor = None, freqs_cis: torch.Tensor = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None) -> torch.Tensor
The [Flux2Transformer2DModel] forward method.
Parameters

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.

Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor = None,
    timestep: torch.LongTensor = None,
    img_ids: torch.Tensor = None,
    txt_ids: torch.Tensor = None,
    guidance: torch.Tensor = None,
    freqs_cis: torch.Tensor = None,
    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
) -> torch.Tensor:
    """
    The [`Flux2Transformer2DModel`] 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
        &#91;diffusers.models.attention_processor&#93;(https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).

&quot;&quot;&quot;
# 0. Handle input arguments
joint_attention_kwargs = joint_attention_kwargs.copy() if joint_attention_kwargs is not None else &#123;&#125;

# Pipeline passes prompt_embeds as a list (one per text encoder); use first
if isinstance(encoder_hidden_states, (list, tuple)):
    if len(encoder_hidden_states) &gt; 1:
        logger.warning(
            &quot;Flux 2 received %d encoder outputs but uses only the first; &quot;
            &quot;indices &gt;= 1 are dropped. Wire multi-encoder support or pass only one.&quot;,
            len(encoder_hidden_states),
        )
    encoder_hidden_states = encoder_hidden_states&#91;0&#93;

num_txt_tokens = encoder_hidden_states.shape&#91;1&#93;

# 0.5. Reshape 5D latent (B, C, T, H, W) from pipeline to 3D (B, T*H*W, C) for transformer
input_was_5d = False
if hidden_states.dim() == 5:
    input_was_5d = True
    b, c_in, t, h, w = hidden_states.shape
    hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape(b, t * h * w, c_in)

# 1. Calculate timestep embedding and modulation parameters (match diffusers: scale to &#91;0, 1000&#93;)
timestep = timestep.to(hidden_states.dtype) * 1000
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)&#91;0&#93;

# 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)

# 3. Compute RoPE positional embeddings when not provided externally
if freqs_cis is None:
    if txt_ids is None or img_ids is None:
        if input_was_5d:
            img_h, img_w = h, w
        else:
            raise ValueError(&quot;Flux 2 4D input requires explicit txt_ids and img_ids &quot;
                             &quot;with real image dimensions; square-image fallback removed.&quot;)

        txt_len = encoder_hidden_states.shape&#91;1&#93;
        txt_ids = torch.cartesian_prod(
            torch.arange(1),
            torch.arange(1),
            torch.arange(1),
            torch.arange(txt_len),
        ).to(device=hidden_states.device)
        img_ids = torch.cartesian_prod(
            torch.arange(1),
            torch.arange(img_h),
            torch.arange(img_w),
            torch.arange(1),
        ).to(device=hidden_states.device)
    freqs_cis = compute_flux2_freqs_cis_from_ids(
        self.rotary_emb,
        txt_ids,
        img_ids,
        device=hidden_states.device,
    )

# 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,
    )
# Concatenate text and image streams for single-block inference
hidden_states = torch.cat(&#91;encoder_hidden_states, hidden_states&#93;, dim=1)

# 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=freqs_cis,
        joint_attention_kwargs=joint_attention_kwargs,
    )
# Remove text tokens from concatenated stream
hidden_states = hidden_states&#91;:, num_txt_tokens:, ...&#93;

# 6. Output layers
hidden_states = self.norm_out(hidden_states, temb)
output, _ = self.proj_out(hidden_states)

# Reshape 3D (B, T*H*W, C) back to 5D (B, C, T, H, W) for scheduler
if input_was_5d:
    output = output.reshape(b, t, h, w, self.out_channels).permute(0, 4, 1, 2, 3)

return output</code></pre></details>