--- title: flux ---

fastvideo.models.dits.flux.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.get_1d_rotary_pos_embed

get_1d_rotary_pos_embed(dim: int, pos: torch.FloatTensor | int, theta: float = 10000.0, theta_rescale_factor: float = 1.0, interpolation_factor: float = 1.0, dtype: torch.dtype = torch.float32, use_real: bool = True, freqs_dtype: torch.dtype | None = None) -> tuple[torch.Tensor, torch.Tensor]

Precompute the frequency tensor for complex exponential (cis) with given dimensions. (Note: cis means cos + i * sin, where i is the imaginary unit.)

This function calculates a frequency tensor with complex exponential using the given dimension 'dim' and the end index 'end'. The 'theta' parameter scales the frequencies.

Parameters

dim: int

Dimension of the frequency tensor.

pos: int or torch.FloatTensor

Position indices for the frequency tensor. [S] or scalar

theta: float

Scaling factor for frequency computation. Defaults to 10000.0.

theta_rescale_factor: float

Rescale factor for theta. Defaults to 1.0.

interpolation_factor: float

Factor to scale positions. Defaults to 1.0.

use_real: bool

If True, output full head_dim with repeated cos/sin for rotate_half style RoPE. If False, output half head_dim for complex style. Defaults to True.

Returns

tuple[torch.Tensor, torch.Tensor]

freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. Shape is [S, D] if use_real=True, [S, D/2] if use_real=False.

fastvideo.models.dits.flux.DistributedAttention

DistributedAttention(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, prefix: str = '', **extra_impl_args)

Distributed attention layer.

fastvideo.models.dits.flux.DistributedAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, original_seq_len: int | None = None, replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, attn_metadata: AttentionMetadata | None = None) -> tuple[torch.Tensor, torch.Tensor | None]

Forward pass for distributed attention.

Parameters

q: torch.Tensor

Query tensor [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor [batch_size, seq_len, num_heads, head_dim]

original_seq_len: int

Original (unpadded) full sequence length

replicated_q: Optional[torch.Tensor]

Replicated query tensor, typically for text tokens

replicated_k: Optional[torch.Tensor]

Replicated key tensor

replicated_v: Optional[torch.Tensor]

Replicated value tensor

Returns

tuple[torch.Tensor, torch.Tensor | None]

Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided

fastvideo.models.dits.flux.DiTConfig

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

fastvideo.models.dits.flux.DiTConfig.add_cli_args

add_cli_args(parser: Any, prefix: str = 'dit-config') -> Any

Add CLI arguments for DiTConfig fields

fastvideo.models.dits.flux.get_forward_context

get_forward_context() -> ForwardContext

Get the current forward context.

fastvideo.models.dits.flux.set_forward_context

set_forward_context(current_timestep, attn_metadata, forward_batch: Optional[ForwardBatch] = None)

A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.

fastvideo.models.dits.flux.ReplicatedLinear

ReplicatedLinear(input_size: int, output_size: int, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Replicated linear layer.

Parameters

input_size: int

input dimension of the linear layer.

output_size: int

output dimension of the linear layer.

bias: bool

If true, add bias.

skip_bias_add: bool

If true, skip adding bias but instead return it.

params_dtype: torch.dtype | None

Data type for the parameters.

quant_config: QuantizationConfig | None

Quantization configure.

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.ReplicatedLinear.weight_loader

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

fastvideo.models.dits.flux.ReplicatedLinear.forward

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

fastvideo.models.dits.flux.ReplicatedLinear.extra_repr

extra_repr() -> str

fastvideo.models.dits.flux.ReplicatedLinear.get_shape_mapping

get_shape_mapping() -> dict

Get the mapping from (input_shape, output_shape) to layer types.

fastvideo.models.dits.flux.ReplicatedLinear.reset_shape_tracking

reset_shape_tracking() -> None

Clear tracked shapes and layer type mappings.

fastvideo.models.dits.flux.ReplicatedLinear.print_shape_summary

print_shape_summary() -> None

Log a summary of all unique shapes and their layer types.

fastvideo.models.dits.flux.Timesteps

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

fastvideo.models.dits.flux.Timesteps.forward

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

fastvideo.models.dits.flux.BaseDiT

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

fastvideo.models.dits.flux.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.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.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.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.CombinedTimestepTextProjEmbeddings

CombinedTimestepTextProjEmbeddings(embedding_dim: int, pooled_projection_dim: int)

fastvideo.models.dits.flux.CombinedTimestepTextProjEmbeddings.forward

forward(timestep: torch.Tensor, pooled_projection: torch.Tensor) -> torch.Tensor

fastvideo.models.dits.flux.SD3AdaLayerNormContinuous

SD3AdaLayerNormContinuous(embedding_dim: int, conditioning_embedding_dim: int, elementwise_affine: bool = True, eps: float = 1e-05, bias: bool = True, norm_type: str = 'layer_norm')

fastvideo.models.dits.flux.SD3AdaLayerNormContinuous.forward

forward(x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor

fastvideo.models.dits.flux.SD3AdaLayerNormZero

SD3AdaLayerNormZero(embedding_dim: int, norm_type: str = 'layer_norm', bias: bool = True)

fastvideo.models.dits.flux.SD3AdaLayerNormZero.forward

forward(x: torch.Tensor, emb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

fastvideo.models.dits.flux.SD3FeedForward

SD3FeedForward(dim: int, dim_out: int | None = None, mult: int = 4, dropout: float = 0.0, activation_fn: str = 'geglu', final_dropout: bool = False, inner_dim: int | None = None, bias: bool = True)

fastvideo.models.dits.flux.SD3FeedForward.forward

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

fastvideo.models.dits.flux.SD3TextProjection

SD3TextProjection(in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = 'silu')

fastvideo.models.dits.flux.SD3TextProjection.forward

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

fastvideo.models.dits.flux.SD3TimestepEmbedding

SD3TimestepEmbedding(in_channels: int, time_embed_dim: int, act_fn: str = 'silu')

fastvideo.models.dits.flux.SD3TimestepEmbedding.forward

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

fastvideo.models.dits.flux.AttentionBackendEnum

AttentionBackendEnum()

fastvideo.models.dits.flux.FluxTransformer2DModelOutput

FluxTransformer2DModelOutput(sample: torch.Tensor)
Source
@dataclass
class FluxTransformer2DModelOutput:
    sample: torch.Tensor

fastvideo.models.dits.flux.FluxPosEmbed

FluxPosEmbed(theta: int, axes_dim: list[int])

1D RoPE axes concatenated per Diffusers FluxPosEmbed.

Source
class FluxPosEmbed(nn.Module):
    """1D RoPE axes concatenated per Diffusers `FluxPosEmbed`."""

    def __init__(self, theta: int, axes_dim: list[int]) -> None:
        super().__init__()
        self.theta = theta
        self.axes_dim = axes_dim

    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        n_axes = ids.shape[-1]
        cos_out: list[torch.Tensor] = []
        sin_out: list[torch.Tensor] = []
        pos = ids.float()
        is_mps = ids.device.type == "mps"
        is_npu = ids.device.type == "npu"
        freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
        for i in range(n_axes):
            cos, sin = get_1d_rotary_pos_embed(
                self.axes_dim[i],
                pos[:, i],
                theta=self.theta,
                use_real=True,
                freqs_dtype=freqs_dtype,
            )
            cos_out.append(cos)
            sin_out.append(sin)
        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
        return freqs_cos, freqs_sin

fastvideo.models.dits.flux.FluxPosEmbed.forward

forward(ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    n_axes = ids.shape[-1]
    cos_out: list[torch.Tensor] = []
    sin_out: list[torch.Tensor] = []
    pos = ids.float()
    is_mps = ids.device.type == "mps"
    is_npu = ids.device.type == "npu"
    freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
    for i in range(n_axes):
        cos, sin = get_1d_rotary_pos_embed(
            self.axes_dim[i],
            pos[:, i],
            theta=self.theta,
            use_real=True,
            freqs_dtype=freqs_dtype,
        )
        cos_out.append(cos)
        sin_out.append(sin)
    freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
    freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
    return freqs_cos, freqs_sin

fastvideo.models.dits.flux.FluxCombinedTimestepGuidanceTextProjEmbeddings

FluxCombinedTimestepGuidanceTextProjEmbeddings(embedding_dim: int, pooled_projection_dim: int)
Source
class FluxCombinedTimestepGuidanceTextProjEmbeddings(nn.Module):

    def __init__(self, embedding_dim: int, pooled_projection_dim: int) -> None:
        super().__init__()
        self.time_proj = Timesteps(
            num_channels=256,
            flip_sin_to_cos=True,
            downscale_freq_shift=0,
        )
        self.timestep_embedder = SD3TimestepEmbedding(
            in_channels=256,
            time_embed_dim=embedding_dim,
            act_fn="silu",
        )
        self.guidance_embedder = SD3TimestepEmbedding(
            in_channels=256,
            time_embed_dim=embedding_dim,
            act_fn="silu",
        )
        self.text_embedder = SD3TextProjection(
            pooled_projection_dim,
            embedding_dim,
            act_fn="silu",
        )

    def forward(
        self,
        timestep: torch.Tensor,
        guidance: torch.Tensor,
        pooled_projection: torch.Tensor,
    ) -> torch.Tensor:
        timesteps_proj = self.time_proj(timestep)
        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
        guidance_proj = self.time_proj(guidance)
        guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
        time_guidance_emb = timesteps_emb + guidance_emb
        pooled_projections = self.text_embedder(pooled_projection)
        return time_guidance_emb + pooled_projections

fastvideo.models.dits.flux.FluxCombinedTimestepGuidanceTextProjEmbeddings.forward

forward(timestep: torch.Tensor, guidance: torch.Tensor, pooled_projection: torch.Tensor) -> torch.Tensor
Source
def forward(
    self,
    timestep: torch.Tensor,
    guidance: torch.Tensor,
    pooled_projection: torch.Tensor,
) -> torch.Tensor:
    timesteps_proj = self.time_proj(timestep)
    timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
    guidance_proj = self.time_proj(guidance)
    guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
    time_guidance_emb = timesteps_emb + guidance_emb
    pooled_projections = self.text_embedder(pooled_projection)
    return time_guidance_emb + pooled_projections

fastvideo.models.dits.flux.FluxAdaLayerNormZeroSingle

FluxAdaLayerNormZeroSingle(embedding_dim: int, bias: bool = True)
Source
class FluxAdaLayerNormZeroSingle(nn.Module):

    def __init__(self, embedding_dim: int, bias: bool = True) -> None:
        super().__init__()
        self.silu = nn.SiLU()
        self.linear = ReplicatedLinear(embedding_dim, 3 * embedding_dim, bias=bias)
        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)

    def forward(
        self,
        x: torch.Tensor,
        emb: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        emb, _ = self.linear(self.silu(emb))
        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
        return x, gate_msa

fastvideo.models.dits.flux.FluxAdaLayerNormZeroSingle.forward

forward(x: torch.Tensor, emb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(
    self,
    x: torch.Tensor,
    emb: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    emb, _ = self.linear(self.silu(emb))
    shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
    x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
    return x, gate_msa

fastvideo.models.dits.flux.FluxJointAttention

FluxJointAttention(dim: int, num_attention_heads: int, attention_head_dim: int, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None)

Joint attention: text tokens precede image tokens (Diffusers order).

Source
class FluxJointAttention(nn.Module):
    """Joint attention: text tokens precede image tokens (Diffusers order)."""

    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
    ) -> None:
        super().__init__()
        self.heads = num_attention_heads
        self.head_dim = attention_head_dim
        self.inner_dim = num_attention_heads * attention_head_dim

        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
        self.norm_added_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
        self.norm_added_k = nn.RMSNorm(attention_head_dim, eps=1e-6)

        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
        self.add_q_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
        self.add_k_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
        self.add_v_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)

        self.to_out = nn.ModuleList([
            ReplicatedLinear(self.inner_dim, dim, bias=True),
            nn.Dropout(0.0),
        ])
        self.to_add_out = ReplicatedLinear(self.inner_dim, dim, bias=True)

        self.attn = DistributedAttention(
            num_heads=num_attention_heads,
            head_size=attention_head_dim,
            causal=False,
            supported_attention_backends=supported_attention_backends,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch_size = hidden_states.shape[0]
        text_seq_len = encoder_hidden_states.shape[1]
        img_seq_len = hidden_states.shape[1]

        q, _ = self.to_q(hidden_states)
        k, _ = self.to_k(hidden_states)
        v, _ = self.to_v(hidden_states)
        q = q.view(batch_size, img_seq_len, self.heads, self.head_dim)
        k = k.view(batch_size, img_seq_len, self.heads, self.head_dim)
        v = v.view(batch_size, img_seq_len, self.heads, self.head_dim)
        q = self.norm_q(q)
        k = self.norm_k(k)

        enc_q, _ = self.add_q_proj(encoder_hidden_states)
        enc_k, _ = self.add_k_proj(encoder_hidden_states)
        enc_v, _ = self.add_v_proj(encoder_hidden_states)
        enc_q = enc_q.view(batch_size, text_seq_len, self.heads, self.head_dim)
        enc_k = enc_k.view(batch_size, text_seq_len, self.heads, self.head_dim)
        enc_v = enc_v.view(batch_size, text_seq_len, self.heads, self.head_dim)
        enc_q = self.norm_added_q(enc_q)
        enc_k = self.norm_added_k(enc_k)

        q = torch.cat([enc_q, q], dim=1)
        k = torch.cat([enc_k, k], dim=1)
        v = torch.cat([enc_v, v], dim=1)

        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)

        joint_out, _ = self.attn(q, k, v)
        joint_out = joint_out.reshape(batch_size, text_seq_len + img_seq_len, self.inner_dim)

        enc_out = joint_out[:, :text_seq_len]
        img_out = joint_out[:, text_seq_len:]

        img_out, _ = self.to_out[0](img_out)
        img_out = self.to_out[1](img_out)
        enc_out, _ = self.to_add_out(enc_out)
        return img_out, enc_out

fastvideo.models.dits.flux.FluxJointAttention.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, image_rotary_emb: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
) -> tuple[torch.Tensor, torch.Tensor]:
    batch_size = hidden_states.shape[0]
    text_seq_len = encoder_hidden_states.shape[1]
    img_seq_len = hidden_states.shape[1]

    q, _ = self.to_q(hidden_states)
    k, _ = self.to_k(hidden_states)
    v, _ = self.to_v(hidden_states)
    q = q.view(batch_size, img_seq_len, self.heads, self.head_dim)
    k = k.view(batch_size, img_seq_len, self.heads, self.head_dim)
    v = v.view(batch_size, img_seq_len, self.heads, self.head_dim)
    q = self.norm_q(q)
    k = self.norm_k(k)

    enc_q, _ = self.add_q_proj(encoder_hidden_states)
    enc_k, _ = self.add_k_proj(encoder_hidden_states)
    enc_v, _ = self.add_v_proj(encoder_hidden_states)
    enc_q = enc_q.view(batch_size, text_seq_len, self.heads, self.head_dim)
    enc_k = enc_k.view(batch_size, text_seq_len, self.heads, self.head_dim)
    enc_v = enc_v.view(batch_size, text_seq_len, self.heads, self.head_dim)
    enc_q = self.norm_added_q(enc_q)
    enc_k = self.norm_added_k(enc_k)

    q = torch.cat([enc_q, q], dim=1)
    k = torch.cat([enc_k, k], dim=1)
    v = torch.cat([enc_v, v], dim=1)

    q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
    k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)

    joint_out, _ = self.attn(q, k, v)
    joint_out = joint_out.reshape(batch_size, text_seq_len + img_seq_len, self.inner_dim)

    enc_out = joint_out[:, :text_seq_len]
    img_out = joint_out[:, text_seq_len:]

    img_out, _ = self.to_out[0](img_out)
    img_out = self.to_out[1](img_out)
    enc_out, _ = self.to_add_out(enc_out)
    return img_out, enc_out

fastvideo.models.dits.flux.FluxSingleStreamAttention

FluxSingleStreamAttention(dim: int, num_attention_heads: int, attention_head_dim: int, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None)

Self-attention on concatenated text+image sequence (single blocks).

Source
class FluxSingleStreamAttention(nn.Module):
    """Self-attention on concatenated text+image sequence (single blocks)."""

    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
    ) -> None:
        super().__init__()
        self.heads = num_attention_heads
        self.head_dim = attention_head_dim
        self.inner_dim = num_attention_heads * attention_head_dim

        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
        self.attn = DistributedAttention(
            num_heads=num_attention_heads,
            head_size=attention_head_dim,
            causal=False,
            supported_attention_backends=supported_attention_backends,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.shape
        q, _ = self.to_q(hidden_states)
        k, _ = self.to_k(hidden_states)
        v, _ = self.to_v(hidden_states)
        q = q.view(batch_size, seq_len, self.heads, self.head_dim)
        k = k.view(batch_size, seq_len, self.heads, self.head_dim)
        v = v.view(batch_size, seq_len, self.heads, self.head_dim)
        q = self.norm_q(q)
        k = self.norm_k(k)
        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
        out, _ = self.attn(q, k, v)
        return out.reshape(batch_size, seq_len, self.inner_dim)

fastvideo.models.dits.flux.FluxSingleStreamAttention.forward

forward(hidden_states: torch.Tensor, image_rotary_emb: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
    batch_size, seq_len, _ = hidden_states.shape
    q, _ = self.to_q(hidden_states)
    k, _ = self.to_k(hidden_states)
    v, _ = self.to_v(hidden_states)
    q = q.view(batch_size, seq_len, self.heads, self.head_dim)
    k = k.view(batch_size, seq_len, self.heads, self.head_dim)
    v = v.view(batch_size, seq_len, self.heads, self.head_dim)
    q = self.norm_q(q)
    k = self.norm_k(k)
    q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
    k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
    out, _ = self.attn(q, k, v)
    return out.reshape(batch_size, seq_len, self.inner_dim)

fastvideo.models.dits.flux.FluxTransformerBlock

FluxTransformerBlock(dim: int, num_attention_heads: int, attention_head_dim: int, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None)
Source
class FluxTransformerBlock(nn.Module):

    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
    ) -> None:
        super().__init__()
        self.norm1 = SD3AdaLayerNormZero(dim)
        self.norm1_context = SD3AdaLayerNormZero(dim)
        self.attn = FluxJointAttention(
            dim=dim,
            num_attention_heads=num_attention_heads,
            attention_head_dim=attention_head_dim,
            supported_attention_backends=supported_attention_backends,
        )
        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.ff = SD3FeedForward(
            dim=dim,
            dim_out=dim,
            activation_fn="gelu-approximate",
        )
        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.ff_context = SD3FeedForward(
            dim=dim,
            dim_out=dim,
            activation_fn="gelu-approximate",
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
        joint_attention_kwargs: dict[str, Any] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        del joint_attention_kwargs
        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
        (norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp,
         c_gate_mlp) = self.norm1_context(encoder_hidden_states, emb=temb)

        attn_output, context_attn_output = self.attn(
            hidden_states=norm_hidden_states,
            encoder_hidden_states=norm_encoder_hidden_states,
            image_rotary_emb=image_rotary_emb,
        )

        attn_output = gate_msa.unsqueeze(1) * attn_output
        hidden_states = hidden_states + attn_output

        norm_hidden_states = self.norm2(hidden_states)
        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
        ff_output = self.ff(norm_hidden_states)
        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output

        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
        encoder_hidden_states = encoder_hidden_states + context_attn_output

        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
        context_ff_output = self.ff_context(norm_encoder_hidden_states)
        encoder_hidden_states = encoder_hidden_states + (c_gate_mlp.unsqueeze(1) * context_ff_output)
        if encoder_hidden_states.dtype == torch.float16:
            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)

        return encoder_hidden_states, hidden_states

fastvideo.models.dits.flux.FluxTransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: tuple[torch.Tensor, torch.Tensor], joint_attention_kwargs: dict[str, Any] | None = None) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    temb: torch.Tensor,
    image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
    joint_attention_kwargs: dict[str, Any] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    del joint_attention_kwargs
    norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
    (norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp,
     c_gate_mlp) = self.norm1_context(encoder_hidden_states, emb=temb)

    attn_output, context_attn_output = self.attn(
        hidden_states=norm_hidden_states,
        encoder_hidden_states=norm_encoder_hidden_states,
        image_rotary_emb=image_rotary_emb,
    )

    attn_output = gate_msa.unsqueeze(1) * attn_output
    hidden_states = hidden_states + attn_output

    norm_hidden_states = self.norm2(hidden_states)
    norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
    ff_output = self.ff(norm_hidden_states)
    hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output

    context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
    encoder_hidden_states = encoder_hidden_states + context_attn_output

    norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
    norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
    context_ff_output = self.ff_context(norm_encoder_hidden_states)
    encoder_hidden_states = encoder_hidden_states + (c_gate_mlp.unsqueeze(1) * context_ff_output)
    if encoder_hidden_states.dtype == torch.float16:
        encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)

    return encoder_hidden_states, hidden_states

fastvideo.models.dits.flux.FluxSingleTransformerBlock

FluxSingleTransformerBlock(dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None)
Source
class FluxSingleTransformerBlock(nn.Module):

    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        mlp_ratio: float = 4.0,
        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
    ) -> None:
        super().__init__()
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.norm = FluxAdaLayerNormZeroSingle(dim)
        self.proj_mlp = ReplicatedLinear(dim, mlp_hidden_dim, bias=True)
        self.act_mlp = nn.GELU(approximate="tanh")
        self.proj_out = ReplicatedLinear(dim + mlp_hidden_dim, dim, bias=True)
        self.attn = FluxSingleStreamAttention(
            dim=dim,
            num_attention_heads=num_attention_heads,
            attention_head_dim=attention_head_dim,
            supported_attention_backends=supported_attention_backends,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
        joint_attention_kwargs: dict[str, Any] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        del joint_attention_kwargs
        text_seq_len = encoder_hidden_states.shape[1]
        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
        residual = hidden_states
        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)[0])
        attn_output = self.attn(
            hidden_states=norm_hidden_states,
            image_rotary_emb=image_rotary_emb,
        )
        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
        gate = gate.unsqueeze(1)
        hidden_states = gate * self.proj_out(hidden_states)[0]
        hidden_states = residual + hidden_states
        if hidden_states.dtype == torch.float16:
            hidden_states = hidden_states.clip(-65504, 65504)
        encoder_hidden_states = hidden_states[:, :text_seq_len]
        hidden_states = hidden_states[:, text_seq_len:]
        return encoder_hidden_states, hidden_states

fastvideo.models.dits.flux.FluxSingleTransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: tuple[torch.Tensor, torch.Tensor], joint_attention_kwargs: dict[str, Any] | None = None) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    temb: torch.Tensor,
    image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
    joint_attention_kwargs: dict[str, Any] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    del joint_attention_kwargs
    text_seq_len = encoder_hidden_states.shape[1]
    hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
    residual = hidden_states
    norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
    mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)[0])
    attn_output = self.attn(
        hidden_states=norm_hidden_states,
        image_rotary_emb=image_rotary_emb,
    )
    hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
    gate = gate.unsqueeze(1)
    hidden_states = gate * self.proj_out(hidden_states)[0]
    hidden_states = residual + hidden_states
    if hidden_states.dtype == torch.float16:
        hidden_states = hidden_states.clip(-65504, 65504)
    encoder_hidden_states = hidden_states[:, :text_seq_len]
    hidden_states = hidden_states[:, text_seq_len:]
    return encoder_hidden_states, hidden_states

fastvideo.models.dits.flux.FluxTransformer2DModel

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

FastVideo FLUX transformer; load Diffusers FLUX safetensors 1:1.

Source
class FluxTransformer2DModel(BaseDiT):
    """FastVideo FLUX transformer; load Diffusers FLUX safetensors 1:1."""

    _fsdp_shard_conditions = [
        lambda n, m: (n.startswith("transformer_blocks.") or n.startswith("single_transformer_blocks.")) and n.split(
            ".")[-1].isdigit(),
    ]
    _compile_conditions = _fsdp_shard_conditions
    # HF weight names already match this module layout (cf. SGLang regex maps).
    param_names_mapping: dict[str, Any] = {}
    reverse_param_names_mapping: dict[str, Any] = {}
    lora_param_names_mapping: dict[str, Any] = {}
    _supported_attention_backends = (
        AttentionBackendEnum.FLASH_ATTN,
        AttentionBackendEnum.TORCH_SDPA,
    )

    def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None:
        del kwargs
        super().__init__(config=config, hf_config=hf_config)
        self.fastvideo_config = config
        self.hf_config = hf_config
        arch = config.arch_config

        out_ch = arch.out_channels
        self.out_channels = out_ch if out_ch is not None else arch.in_channels
        self.inner_dim = arch.num_attention_heads * arch.attention_head_dim
        self.hidden_size = self.inner_dim
        self.num_attention_heads = arch.num_attention_heads
        self.num_channels_latents = arch.in_channels

        axes_list = list(arch.axes_dims_rope)
        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_list)
        if arch.guidance_embeds:
            self.time_text_embed = FluxCombinedTimestepGuidanceTextProjEmbeddings(
                embedding_dim=self.inner_dim,
                pooled_projection_dim=arch.pooled_projection_dim,
            )
        else:
            self.time_text_embed = CombinedTimestepTextProjEmbeddings(
                embedding_dim=self.inner_dim,
                pooled_projection_dim=arch.pooled_projection_dim,
            )
        self.context_embedder = ReplicatedLinear(arch.joint_attention_dim, self.inner_dim)
        self.x_embedder = ReplicatedLinear(arch.in_channels, self.inner_dim)

        self.transformer_blocks = nn.ModuleList([
            FluxTransformerBlock(
                dim=self.inner_dim,
                num_attention_heads=arch.num_attention_heads,
                attention_head_dim=arch.attention_head_dim,
                supported_attention_backends=self._supported_attention_backends,
            ) for _ in range(arch.num_layers)
        ])
        self.single_transformer_blocks = nn.ModuleList([
            FluxSingleTransformerBlock(
                dim=self.inner_dim,
                num_attention_heads=arch.num_attention_heads,
                attention_head_dim=arch.attention_head_dim,
                supported_attention_backends=self._supported_attention_backends,
            ) for _ in range(arch.num_single_layers)
        ])

        self.norm_out = SD3AdaLayerNormContinuous(
            self.inner_dim,
            self.inner_dim,
            elementwise_affine=False,
            eps=1e-6,
            bias=True,
            norm_type="layer_norm",
        )
        self.proj_out = ReplicatedLinear(
            self.inner_dim,
            arch.patch_size * arch.patch_size * self.out_channels,
            bias=True,
        )
        self.gradient_checkpointing = False
        self.__post_init__()

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor | None = None,
        pooled_projections: torch.Tensor | None = None,
        timestep: torch.LongTensor | torch.Tensor | None = None,
        img_ids: torch.Tensor | None = None,
        txt_ids: torch.Tensor | None = None,
        guidance: torch.Tensor | None = None,
        joint_attention_kwargs: dict[str, Any] | None = None,
        return_dict: bool = True,
        controlnet_block_samples: Any | None = None,
        controlnet_single_block_samples: Any | None = None,
        controlnet_blocks_repeat: bool = False,
        **kwargs: Any,
    ) -> FluxTransformer2DModelOutput | tuple[torch.Tensor, ...]:
        del kwargs
        if encoder_hidden_states is None:
            raise ValueError("encoder_hidden_states must be provided")
        if pooled_projections is None:
            raise ValueError("pooled_projections must be provided")
        if timestep is None:
            raise ValueError("timestep must be provided")
        if img_ids is None or txt_ids is None:
            raise ValueError("img_ids and txt_ids must be provided")

        arch = self.fastvideo_config.arch_config
        if arch.guidance_embeds and guidance is None:
            raise ValueError("guidance must be provided when guidance_embeds=True")

        if timestep.dim() == 0:
            timestep = timestep[None]
        if timestep.dim() > 1:
            timestep = timestep.reshape(-1)
        if timestep.shape[0] == 1 and hidden_states.shape[0] > 1:
            timestep = timestep.expand(hidden_states.shape[0])

        try:
            get_forward_context()
            forward_context = nullcontext()
        except AssertionError:
            if timestep.numel() == 0:
                ts0 = 0
            elif torch.is_floating_point(timestep):
                ts0 = int(round(timestep[0].item() * 1000))
            else:
                ts0 = int(timestep[0].item())
            forward_context = set_forward_context(current_timestep=ts0, attn_metadata=None)

        with forward_context:
            hidden_states, _ = self.x_embedder(hidden_states)

            ts = timestep.to(hidden_states.dtype) * 1000
            g = None if guidance is None else guidance.to(hidden_states.dtype) * 1000

            if arch.guidance_embeds:
                assert g is not None
                temb = self.time_text_embed(ts, g, pooled_projections)
            else:
                temb = self.time_text_embed(timestep=ts, pooled_projection=pooled_projections)

            encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)

            if txt_ids.ndim == 3:
                txt_ids = txt_ids[0]
            if img_ids.ndim == 3:
                img_ids = img_ids[0]

            ids = torch.cat((txt_ids, img_ids), dim=0)
            image_rotary_emb = self.pos_embed(ids)

            jkwargs = joint_attention_kwargs or {}

            for idx, block in enumerate(self.transformer_blocks):
                encoder_hidden_states, hidden_states = block(
                    hidden_states=hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    temb=temb,
                    image_rotary_emb=image_rotary_emb,
                    joint_attention_kwargs=jkwargs,
                )
                if controlnet_block_samples:
                    interval = len(self.transformer_blocks) / len(controlnet_block_samples)
                    interval = int(math.ceil(interval))
                    if controlnet_blocks_repeat:
                        hidden_states = hidden_states + controlnet_block_samples[idx % len(controlnet_block_samples)]
                    else:
                        hidden_states = hidden_states + controlnet_block_samples[idx // interval]

            for idx, block in enumerate(self.single_transformer_blocks):
                encoder_hidden_states, hidden_states = block(
                    hidden_states=hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    temb=temb,
                    image_rotary_emb=image_rotary_emb,
                    joint_attention_kwargs=jkwargs,
                )
                if controlnet_single_block_samples:
                    interval = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
                    interval = int(math.ceil(interval))
                    hidden_states = hidden_states + controlnet_single_block_samples[idx // interval]

            hidden_states = self.norm_out(hidden_states, temb)
            output, _ = self.proj_out(hidden_states)

            if not return_dict:
                return (output, )
            return FluxTransformer2DModelOutput(sample=output)

fastvideo.models.dits.flux.FluxTransformer2DModel.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, pooled_projections: torch.Tensor | None = None, timestep: torch.LongTensor | torch.Tensor | None = None, img_ids: torch.Tensor | None = None, txt_ids: torch.Tensor | None = None, guidance: torch.Tensor | None = None, joint_attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, controlnet_block_samples: Any | None = None, controlnet_single_block_samples: Any | None = None, controlnet_blocks_repeat: bool = False, **kwargs: Any) -> FluxTransformer2DModelOutput | tuple[torch.Tensor, ...]
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor | None = None,
    pooled_projections: torch.Tensor | None = None,
    timestep: torch.LongTensor | torch.Tensor | None = None,
    img_ids: torch.Tensor | None = None,
    txt_ids: torch.Tensor | None = None,
    guidance: torch.Tensor | None = None,
    joint_attention_kwargs: dict[str, Any] | None = None,
    return_dict: bool = True,
    controlnet_block_samples: Any | None = None,
    controlnet_single_block_samples: Any | None = None,
    controlnet_blocks_repeat: bool = False,
    **kwargs: Any,
) -> FluxTransformer2DModelOutput | tuple[torch.Tensor, ...]:
    del kwargs
    if encoder_hidden_states is None:
        raise ValueError("encoder_hidden_states must be provided")
    if pooled_projections is None:
        raise ValueError("pooled_projections must be provided")
    if timestep is None:
        raise ValueError("timestep must be provided")
    if img_ids is None or txt_ids is None:
        raise ValueError("img_ids and txt_ids must be provided")

    arch = self.fastvideo_config.arch_config
    if arch.guidance_embeds and guidance is None:
        raise ValueError("guidance must be provided when guidance_embeds=True")

    if timestep.dim() == 0:
        timestep = timestep[None]
    if timestep.dim() > 1:
        timestep = timestep.reshape(-1)
    if timestep.shape[0] == 1 and hidden_states.shape[0] > 1:
        timestep = timestep.expand(hidden_states.shape[0])

    try:
        get_forward_context()
        forward_context = nullcontext()
    except AssertionError:
        if timestep.numel() == 0:
            ts0 = 0
        elif torch.is_floating_point(timestep):
            ts0 = int(round(timestep[0].item() * 1000))
        else:
            ts0 = int(timestep[0].item())
        forward_context = set_forward_context(current_timestep=ts0, attn_metadata=None)

    with forward_context:
        hidden_states, _ = self.x_embedder(hidden_states)

        ts = timestep.to(hidden_states.dtype) * 1000
        g = None if guidance is None else guidance.to(hidden_states.dtype) * 1000

        if arch.guidance_embeds:
            assert g is not None
            temb = self.time_text_embed(ts, g, pooled_projections)
        else:
            temb = self.time_text_embed(timestep=ts, pooled_projection=pooled_projections)

        encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)

        if txt_ids.ndim == 3:
            txt_ids = txt_ids[0]
        if img_ids.ndim == 3:
            img_ids = img_ids[0]

        ids = torch.cat((txt_ids, img_ids), dim=0)
        image_rotary_emb = self.pos_embed(ids)

        jkwargs = joint_attention_kwargs or {}

        for idx, block in enumerate(self.transformer_blocks):
            encoder_hidden_states, hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                temb=temb,
                image_rotary_emb=image_rotary_emb,
                joint_attention_kwargs=jkwargs,
            )
            if controlnet_block_samples:
                interval = len(self.transformer_blocks) / len(controlnet_block_samples)
                interval = int(math.ceil(interval))
                if controlnet_blocks_repeat:
                    hidden_states = hidden_states + controlnet_block_samples[idx % len(controlnet_block_samples)]
                else:
                    hidden_states = hidden_states + controlnet_block_samples[idx // interval]

        for idx, block in enumerate(self.single_transformer_blocks):
            encoder_hidden_states, hidden_states = block(
                hidden_states=hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                temb=temb,
                image_rotary_emb=image_rotary_emb,
                joint_attention_kwargs=jkwargs,
            )
            if controlnet_single_block_samples:
                interval = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
                interval = int(math.ceil(interval))
                hidden_states = hidden_states + controlnet_single_block_samples[idx // interval]

        hidden_states = self.norm_out(hidden_states, temb)
        output, _ = self.proj_out(hidden_states)

        if not return_dict:
            return (output, )
        return FluxTransformer2DModelOutput(sample=output)