flux_2 ¶

Classes¶

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)

Bases: 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.

Source code in fastvideo/models/dits/flux_2.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,
    )

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)

Bases: 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.

Source code in fastvideo/models/dits/flux_2.py

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)

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)

Bases: 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 for a visual depiction of this type of transformer block.

Source code in fastvideo/models/dits/flux_2.py

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

fastvideo.models.dits.flux_2.Flux2PosEmbed ¶

Flux2PosEmbed(theta: int, axes_dim: List[int])

Bases: Module

Flux2 N-D RoPE: matches Diffusers Flux2PosEmbed (transformer_flux2).

Source code in fastvideo/models/dits/flux_2.py

def __init__(self, theta: int, axes_dim: List[int]):
    super().__init__()
    self.theta = theta
    self.axes_dim = axes_dim

Methods:¶

fastvideo.models.dits.flux_2.Flux2PosEmbed.forward ¶

forward(ids: Tensor) -> tuple[Tensor, Tensor]

Per Diffusers: loop axes, 1D RoPE each, concat on channel dim.

Source code in fastvideo/models/dits/flux_2.py

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: Tensor) -> tuple[Tensor, Tensor]

Same as forward; kept for callers that pre-flatten positions.

Source code in fastvideo/models/dits/flux_2.py

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.Flux2SwiGLU ¶

Flux2SwiGLU(compute_in_fp32: bool = False)

Bases: 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.

Source code in fastvideo/models/dits/flux_2.py

def __init__(self, compute_in_fp32: bool = False):
    super().__init__()
    self.gate_fn = nn.SiLU()
    self.compute_in_fp32 = compute_in_fp32

fastvideo.models.dits.flux_2.Flux2Transformer2DModel ¶

Flux2Transformer2DModel(config: Flux2Config, hf_config: dict[str, Any])

Bases: BaseDiT

The Transformer model introduced in Flux 2.

Reference: https://blackforestlabs.ai/announcing-black-forest-labs/

Source code in fastvideo/models/dits/flux_2.py

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"]

Methods:¶

fastvideo.models.dits.flux_2.Flux2Transformer2DModel.forward ¶

forward(hidden_states: Tensor, encoder_hidden_states: Tensor = None, timestep: LongTensor = None, img_ids: Tensor = None, txt_ids: Tensor = None, guidance: Tensor = None, freqs_cis: Tensor = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None) -> Tensor

The [Flux2Transformer2DModel] forward method.

Parameters:

Name	Type	Description	Default
`hidden_states`	`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`	Input `hidden_states`.	required
`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.	`None`
`timestep`	`torch.LongTensor`	Used to indicate denoising step.	`None`
`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.	`None`

Source code in fastvideo/models/dits/flux_2.py

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

Functions:¶

fastvideo.models.dits.flux_2.apply_qk_norm ¶

apply_qk_norm(q: Tensor, k: Tensor, q_norm: Module, k_norm: Module, head_dim: int, allow_inplace: bool = True) -> Tuple[Tensor, Tensor]

Apply QK normalization for query and key tensors.

Source code in fastvideo/models/dits/flux_2.py

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.compute_flux2_freqs_cis_from_ids ¶

compute_flux2_freqs_cis_from_ids(rotary_emb: Flux2PosEmbed, text_ids: Tensor, latent_ids: Tensor, device: device, dtype: dtype | None = None) -> tuple[Tensor, 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 code in fastvideo/models/dits/flux_2.py

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)