--- title: flux_2 ---

fastvideo.models.dits.flux_2.Flux2Config

Flux2Config(arch_config: DiTArchConfig = Flux2ArchConfig(), prefix: str = 'Flux', quant_config: QuantizationConfig | None = None, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

Configuration for Flux2 transformer model.

fastvideo.models.dits.flux_2.LocalAttention

LocalAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, default_backend: AttentionBackendEnum | None = None, **extra_impl_args)

Attention layer.

fastvideo.models.dits.flux_2.LocalAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor

Apply local attention between query, key and value tensors.

Parameters

q: torch.Tensor

Query tensor of shape [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor of shape [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor of shape [batch_size, seq_len, num_heads, head_dim]

Returns

torch.Tensor

torch.Tensor: Output tensor after local attention

fastvideo.models.dits.flux_2.get_tp_world_size

get_tp_world_size() -> int

Return world size for the tensor model parallel group.

fastvideo.models.dits.flux_2.Timesteps

Timesteps(num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1)

fastvideo.models.dits.flux_2.Timesteps.forward

forward(timesteps: torch.Tensor) -> torch.Tensor

fastvideo.models.dits.flux_2.ColumnParallelLinear

ColumnParallelLinear(input_size: int, output_size: int, bias: bool = True, gather_output: bool = False, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, output_sizes: list[int] | None = None, prefix: str = '')

Linear layer with column parallelism.

The linear layer is defined as Y = XA + b. A is parallelized along its second dimension as A = [A_1, ..., A_p].

Parameters

input_size: int

first dimension of matrix A.

output_size: int

second dimension of matrix A.

bias: bool

If true, add bias.

gather_output: bool

If true, call all-gather on output and make Y available to all GPUs, otherwise, every GPU will have its output which is Y_i = XA_i

skip_bias_add: bool

This was added to enable performance optimizations where bias can be fused with other element-wise operations. we skip adding bias but instead return it.

params_dtype: torch.dtype | None

Data type for the parameters.

quant_config: QuantizationConfig | None

Quantization configure.

output_sizes: list[int] | None

list of output sizes packed into one output, like for QKV the list would be size 3.

prefix: str

The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

fastvideo.models.dits.flux_2.ColumnParallelLinear.weight_loader

weight_loader(param: Parameter, loaded_weight: torch.Tensor) -> None

fastvideo.models.dits.flux_2.ColumnParallelLinear.weight_loader_v2

weight_loader_v2(param: Parameter, loaded_weight: torch.Tensor) -> None

fastvideo.models.dits.flux_2.ColumnParallelLinear.forward

forward(input_: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]

fastvideo.models.dits.flux_2.ColumnParallelLinear.extra_repr

extra_repr() -> str

fastvideo.models.dits.flux_2.apply_rotary_emb

apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor | tuple[torch.Tensor, torch.Tensor], use_real: bool = True, use_real_unbind_dim: int = -1, sequence_dim: int = 2) -> torch.Tensor

Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are returned as real tensors. Args: x (torch.Tensor): Query or key tensor to apply rotary embeddings. [B, H, S, D] if sequence_dim=2 else [B, S, H, D]. freqs_cis (Tuple[torch.Tensor]): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) sequence_dim: 1 = sequence at dim 1 (cos [1,S,1,D], x [B,S,H,D]); 2 = sequence at dim 2 (cos [1,1,S,D], x [B,H,S,D]). Returns: Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.

fastvideo.models.dits.flux_2.BaseDiT

BaseDiT(config: DiTConfig, hf_config: dict[str, Any], **kwargs)

fastvideo.models.dits.flux_2.BaseDiT.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance = None, **kwargs) -> torch.Tensor

fastvideo.models.dits.flux_2.BaseDiT.preprocess_loaded_weight

preprocess_loaded_weight(name: str, tensor: torch.Tensor) -> torch.Tensor

Adapt a checkpoint tensor before sharding; default models are unchanged.

fastvideo.models.dits.flux_2.BaseDiT.inference_config_overrides

inference_config_overrides() -> dict[str, Any]

Extra config.json keys a deployable checkpoint of this model needs.

fastvideo.models.dits.flux_2.BaseDiT.inference_export_expected_shape

inference_export_expected_shape(output_key: str, base_shape: tuple[int, ...]) -> tuple[int, ...]

The exported shape of output_key given the base checkpoint's shape.

fastvideo.models.dits.flux_2.AttentionBackendEnum

AttentionBackendEnum()

fastvideo.models.dits.flux_2.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.models.dits.flux_2.TimestepEmbedding

TimestepEmbedding(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)
Source
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

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 = "layer_norm",
    ):
        super().__init__()
        self.silu = nn.SiLU()
        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
        if norm_type != "layer_norm":
            raise ValueError(f"unknown norm_type {norm_type}")
        self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine, bias=bias)

    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.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) -> 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.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``.
    """

    def __init__(
        self,
        dim: int,
        dim_out: Optional[int] = None,
        mult: float = 3.0,
        inner_dim: Optional[int] = 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) -> 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.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``.
    """

    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-6,
        out_dim: int = None,
        elementwise_affine: bool = True,
        supported_attention_backends: Optional[Tuple[AttentionBackendEnum, ...]] = 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([])
        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[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([encoder_query, query], dim=1)
            key = torch.cat([encoder_key, key], dim=1)
            value = torch.cat([encoder_value, value], 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: [B, S, H, D] with sequence_dim=1 (cos/sin [S, D]).
            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(
                [
                    encoder_hidden_states.shape[1],
                    hidden_states.shape[1] - encoder_hidden_states.shape[1],
                ],
                dim=1,
            )
            encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)

        hidden_states, _ = self.to_out[0](hidden_states)
        hidden_states = self.to_out[1](hidden_states)

        if encoder_hidden_states is not None:
            return hidden_states, encoder_hidden_states
        else:
            return hidden_states

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([encoder_query, query], dim=1)
        key = torch.cat([encoder_key, key], dim=1)
        value = torch.cat([encoder_value, value], 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: [B, S, H, D] with sequence_dim=1 (cos/sin [S, D]).
        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(
            [
                encoder_hidden_states.shape[1],
                hidden_states.shape[1] - encoder_hidden_states.shape[1],
            ],
            dim=1,
        )
        encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)

    hidden_states, _ = self.to_out[0](hidden_states)
    hidden_states = self.to_out[1](hidden_states)

    if encoder_hidden_states is not None:
        return hidden_states, encoder_hidden_states
    else:
        return hidden_states

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 [ViT-22B
    paper](https://arxiv.org/abs/2302.05442) for a visual depiction of this type of transformer block.
    """

    # Does not support QKV fusion as the QKV projections are always fused
    _supports_qkv_fusion = False

    def __init__(
        self,
        query_dim: int,
        num_heads: int = 8,
        dim_head: int = 64,
        dropout: float = 0.0,
        bias: bool = False,
        out_bias: bool = True,
        eps: float = 1e-6,
        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,
    ):
        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[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 [B, S, H, D], 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([hidden_states, mlp_hidden_states], dim=-1)
        hidden_states, _ = self.to_out(hidden_states)

        return hidden_states

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 [B, S, H, D], 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([hidden_states, mlp_hidden_states], dim=-1)
    hidden_states, _ = self.to_out(hidden_states)

    return hidden_states

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[Tuple[AttentionBackendEnum, ...]] = 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 "parallel" transformer block. See the [ViT-22B paper](https://arxiv.org/abs/2302.05442)
        # for a visual depiction of this type of transformer block.
        self.attn = Flux2ParallelSelfAttention(
            query_dim=dim,
            dim_head=attention_head_dim,
            num_heads=num_attention_heads,
            out_dim=dim,
            bias=bias,
            out_bias=bias,
            eps=eps,
            mlp_ratio=mlp_ratio,
            mlp_mult_factor=2,
            supported_attention_backends=supported_attention_backends,
        )

    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 {}
        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[:, :text_seq_len],
                hidden_states[:, text_seq_len:],
            )
            return encoder_hidden_states, hidden_states
        else:
            return hidden_states

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 {}
    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[:, :text_seq_len],
            hidden_states[:, text_seq_len:],
        )
        return encoder_hidden_states, hidden_states
    else:
        return hidden_states

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[Tuple[AttentionBackendEnum, ...]] = 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[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: [1, 1, self.dim]
        (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = (temb_mod_params_img)
        (c_shift_msa, c_scale_msa, c_gate_msa), (
            c_shift_mlp,
            c_scale_mlp,
            c_gate_mlp,
        ) = temb_mod_params_txt

        # Img stream
        norm_hidden_states = self.norm1(hidden_states)
        norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa

        # Conditioning txt stream
        norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
        norm_encoder_hidden_states = (1 + c_scale_msa) * norm_encoder_hidden_states + c_shift_msa

        # Attention on concatenated img + txt stream
        attention_outputs = self.attn(
            hidden_states=norm_hidden_states,
            encoder_hidden_states=norm_encoder_hidden_states,
            freqs_cis=freqs_cis,
        )

        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

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: [1, 1, self.dim]
    (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = (temb_mod_params_img)
    (c_shift_msa, c_scale_msa, c_gate_msa), (
        c_shift_mlp,
        c_scale_mlp,
        c_gate_mlp,
    ) = temb_mod_params_txt

    # Img stream
    norm_hidden_states = self.norm1(hidden_states)
    norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa

    # Conditioning txt stream
    norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
    norm_encoder_hidden_states = (1 + c_scale_msa) * norm_encoder_hidden_states + c_shift_msa

    # Attention on concatenated img + txt stream
    attention_outputs = self.attn(
        hidden_states=norm_hidden_states,
        encoder_hidden_states=norm_encoder_hidden_states,
        freqs_cis=freqs_cis,
    )

    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

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[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

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

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) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]:
        mod = self.act_fn(temb)
        mod, _ = self.linear(mod)

        if mod.ndim == 2:
            mod = mod.unsqueeze(1)
        mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1)
        # Return tuple of 3-tuples of modulation params shift/scale/gate
        return tuple(mod_params[3 * i:3 * (i + 1)] for i in range(self.mod_param_sets))

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[3 * i:3 * (i + 1)] for i in range(self.mod_param_sets))

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[int]):
        super().__init__()
        self.theta = theta
        self.axes_dim = axes_dim

    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)

    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.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.
    """
    tid = text_ids
    lid = latent_ids
    if tid.ndim == 3:
        tid = tid[0]
    if lid.ndim == 3:
        lid = lid[0]
    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([text_cos, img_cos], dim=0).to(device=device)
    sin = torch.cat([text_sin, img_sin], dim=0).to(device=device)
    # Match Diffusers: keep RoPE cos/sin in fp32 (``dtype`` unused; callers may still pass it).
    _ = dtype
    return (cos, sin)

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/

    """

    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[str, Any]):
        super().__init__(config=config, hf_config=hf_config)
        patch_size: int = config.patch_size
        in_channels: int = config.in_channels
        out_channels: Optional[int] = config.out_channels
        num_layers: int = config.num_layers
        num_single_layers: int = config.num_single_layers
        attention_head_dim: int = config.attention_head_dim
        num_attention_heads: int = config.num_attention_heads
        joint_attention_dim: int = config.joint_attention_dim
        timestep_guidance_channels: int = config.timestep_guidance_channels
        mlp_ratio: float = config.mlp_ratio
        axes_dims_rope: Tuple[int, ...] = config.axes_dims_rope
        rope_theta: int = config.rope_theta
        eps: float = config.eps
        guidance_embeds: bool = getattr(config, "guidance_embeds", True)
        self.out_channels = out_channels or in_channels
        self.inner_dim = num_attention_heads * attention_head_dim
        self.guidance_embeds = guidance_embeds

        # 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, "ff_context_swiglu_fp32", False)
        self.transformer_blocks = nn.ModuleList([
            Flux2TransformerBlock(
                dim=self.inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                mlp_ratio=mlp_ratio,
                eps=eps,
                bias=False,
                supported_attention_backends=supported_attention_backends,
                ff_context_swiglu_fp32=ff_ctx_swiglu,
            ) for _ in range(num_layers)
        ])

        # 6. Single Stream Transformer Blocks
        self.single_transformer_blocks = nn.ModuleList([
            Flux2SingleTransformerBlock(
                dim=self.inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                mlp_ratio=mlp_ratio,
                eps=eps,
                bias=False,
                supported_attention_backends=supported_attention_backends,
            ) for _ in range(num_single_layers)
        ])

        # 7. Output layers
        self.norm_out = AdaLayerNormContinuous(
            self.inner_dim,
            self.inner_dim,
            elementwise_affine=False,
            eps=eps,
            bias=False,
        )
        self.proj_out = ColumnParallelLinear(
            self.inner_dim,
            patch_size * patch_size * self.out_channels,
            bias=False,
            gather_output=True,
        )

        self.layer_names = ["transformer_blocks", "single_transformer_blocks"]

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor = None,
        timestep: torch.LongTensor = None,
        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
                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).

        """
        # 0. Handle input arguments
        joint_attention_kwargs = joint_attention_kwargs.copy() if joint_attention_kwargs is not None else {}

        # 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) > 1:
                logger.warning(
                    "Flux 2 received %d encoder outputs but uses only the first; "
                    "indices >= 1 are dropped. Wire multi-encoder support or pass only one.",
                    len(encoder_hidden_states),
                )
            encoder_hidden_states = encoder_hidden_states[0]

        num_txt_tokens = encoder_hidden_states.shape[1]

        # 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 [0, 1000])
        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)[0]

        # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states)
        hidden_states, _ = self.x_embedder(hidden_states)
        encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)

        # 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("Flux 2 4D input requires explicit txt_ids and img_ids "
                                     "with real image dimensions; square-image fallback removed.")

                txt_len = encoder_hidden_states.shape[1]
                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([encoder_hidden_states, hidden_states], 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[:, num_txt_tokens:, ...]

        # 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

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
            [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).

    """
    # 0. Handle input arguments
    joint_attention_kwargs = joint_attention_kwargs.copy() if joint_attention_kwargs is not None else {}

    # 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) > 1:
            logger.warning(
                "Flux 2 received %d encoder outputs but uses only the first; "
                "indices >= 1 are dropped. Wire multi-encoder support or pass only one.",
                len(encoder_hidden_states),
            )
        encoder_hidden_states = encoder_hidden_states[0]

    num_txt_tokens = encoder_hidden_states.shape[1]

    # 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 [0, 1000])
    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)[0]

    # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states)
    hidden_states, _ = self.x_embedder(hidden_states)
    encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)

    # 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("Flux 2 4D input requires explicit txt_ids and img_ids "
                                 "with real image dimensions; square-image fallback removed.")

            txt_len = encoder_hidden_states.shape[1]
            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([encoder_hidden_states, hidden_states], 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[:, num_txt_tokens:, ...]

    # 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