Add vendor/mage_flow/models/modules/mage_layers.py

vendor/mage_flow/models/modules/mage_layers.py

@@ -0,0 +1,733 @@
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.attention import FeedForward
+from diffusers.models.embeddings import TimestepEmbedding
+from diffusers.models.normalization import RMSNorm
+from ._attn_backend import flash_attn_varlen_func
+from torch import Tensor
+from torch._dynamo import allow_in_graph as maybe_allow_in_graph
+
+
+def apply_rotary_emb_mageflow(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply complex rotary embeddings to `x` ([B, S, H, D]) using `freqs_cis`
+    (the MageFlowEmbedRope 2D multi-scale RoPE, adjacent-pair complex convention)."""
+    x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    freqs_cis = freqs_cis.unsqueeze(1)
+    x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(-2)
+    return x_out.type_as(x)
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+) -> torch.Tensor:
+    """Sinusoidal timestep embeddings (DDPM convention).
+
+    NOTE: kept vendored (not diffusers') because the frequency table is
+    downcast to ``timesteps.dtype`` (bf16) here — the model was trained with
+    this exact bf16 rounding, so diffusers' fp32 variant produces a slightly
+    different embedding and degrades outputs.
+    """
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent).to(timesteps.dtype)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # scale embeddings
+    emb = scale * emb
+
+    # concat sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # flip sine and cosine embeddings
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # zero pad
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+class Timesteps(nn.Module):
+    def __init__(
+        self,
+        num_channels: int,
+        flip_sin_to_cos: bool,
+        downscale_freq_shift: float,
+        scale: int = 1,
+    ):
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        t_emb = get_timestep_embedding(
+            timesteps,
+            self.num_channels,
+            flip_sin_to_cos=self.flip_sin_to_cos,
+            downscale_freq_shift=self.downscale_freq_shift,
+            scale=self.scale,
+        )
+        return t_emb
+
+
+class MageFlowTimestepProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+    def forward(self, timestep, hidden_states):
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))  # (N, D)
+
+        conditioning = timesteps_emb
+
+        return conditioning
+
+
+class MageFlowEmbedRope(nn.Module):
+    def __init__(self, theta: int, axes_dim: list[int], scale_rope=False):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        pos_index = torch.arange(4096)
+        neg_index = torch.arange(4096).flip(0) * -1 - 1
+        self.pos_freqs = torch.cat(
+            [
+                self.rope_params(pos_index, self.axes_dim[0], self.theta),
+                self.rope_params(pos_index, self.axes_dim[1], self.theta),
+                self.rope_params(pos_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+        self.neg_freqs = torch.cat(
+            [
+                self.rope_params(neg_index, self.axes_dim[0], self.theta),
+                self.rope_params(neg_index, self.axes_dim[1], self.theta),
+                self.rope_params(neg_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+
+        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
+        self.scale_rope = scale_rope
+        self.video_freq_cache = {}
+
+    def rope_params(self, index, dim, theta=10000):
+        """
+        Args:
+            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
+        """
+        assert dim % 2 == 0
+        freqs = torch.outer(
+            index,
+            1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)),
+        )
+        freqs = torch.polar(torch.ones_like(freqs), freqs)
+        return freqs
+
+    def forward(
+        self,
+        video_fhw: tuple[int, int, int] | list[tuple[int, int, int]],
+        device: torch.device,
+        max_img_len: int = None,
+    ) -> torch.Tensor:
+        """Compute the vision RoPE frequencies (`vid_freqs`) for the packed image
+        tokens. Text tokens are NOT rotated, so no text RoPE is computed.
+
+        Args:
+            video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`):
+                A list of 3 integers [frame, height, width] representing the shape of the video.
+            device: (`torch.device`):
+                The device on which to perform the RoPE computation.
+        """
+        if self.pos_freqs.device != device:
+            self.pos_freqs = self.pos_freqs.to(device)
+            self.neg_freqs = self.neg_freqs.to(device)
+
+        if isinstance(video_fhw, list):
+            video_fhw = video_fhw[0]
+        if not isinstance(video_fhw, list):
+            video_fhw = [video_fhw]
+
+        vid_freqs = []
+        for idx, fhw in enumerate(video_fhw):
+            frame, height, width = fhw
+            # RoPE frequencies are cached manually
+            key = (frame, height, width, idx)
+            if key not in self.video_freq_cache:
+                self.video_freq_cache[key] = self._compute_video_freqs(frame, height, width, idx)
+            vid_freqs.append(self.video_freq_cache[key].to(device))
+
+        vid_freqs = torch.cat(vid_freqs, dim=0)
+
+        if max_img_len is not None and vid_freqs.shape[0] < max_img_len:
+            pad_len = max_img_len - vid_freqs.shape[0]
+            vid_freqs = torch.nn.functional.pad(vid_freqs, (0, 0, 0, pad_len))
+
+        return vid_freqs
+
+    def _compute_video_freqs(self, frame: int, height: int, width: int, idx: int = 0) -> torch.Tensor:
+        seq_lens = frame * height * width
+        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+        freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+
+        freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
+        if self.scale_rope:
+            freqs_height = torch.cat(
+                [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
+                dim=0,
+            )
+            freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = torch.cat(
+                [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
+                dim=0,
+            )
+            freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1)
+        else:
+            freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1)
+
+        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
+        return freqs.clone().contiguous()
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        cross_attention_dim: int | None = None,
+        heads: int = 8,
+        kv_heads: int | None = None,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        scale_qk: bool = True,
+        added_kv_proj_dim: int | None = None,
+        added_proj_bias: bool | None = True,
+        out_bias: bool = True,
+        eps: float = 1e-5,
+        processor=None,
+        out_dim: int = None,
+        out_context_dim: int = None,
+        elementwise_affine: bool = True,
+    ):
+        super().__init__()
+        # logger.info(f"processor: {processor}")
+
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+        self.query_dim = query_dim
+        self.use_bias = bias
+        self.is_cross_attention = cross_attention_dim is not None
+        self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+        self.fused_projections = False
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim
+
+        self.scale_qk = scale_qk
+        self.scale = dim_head**-0.5 if self.scale_qk else 1.0
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        # for slice_size > 0 the attention score computation
+        # is split across the batch axis to save memory
+        # You can set_slice_size with `set_attention_slice`
+        self.sliceable_head_dim = heads
+
+        self.added_kv_proj_dim = added_kv_proj_dim
+
+        # qk_norm is always "rms_norm" for MageFlow.
+        self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+        self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+
+        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+        self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+
+        self.added_proj_bias = added_proj_bias
+        if self.added_kv_proj_dim is not None:
+            self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+            self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+            self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+            self.norm_added_q = RMSNorm(dim_head, eps=eps)
+            self.norm_added_k = RMSNorm(dim_head, eps=eps)
+        else:
+            self.add_q_proj = None
+            self.add_k_proj = None
+            self.add_v_proj = None
+            self.norm_added_q = None
+            self.norm_added_k = None
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+        self.to_out.append(nn.Dropout(dropout))
+
+        self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
+
+        self.set_processor(processor)
+
+    def set_processor(self, processor) -> None:
+        self.processor = processor
+
+    def get_processor(self):
+        return self.processor
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        txt_cu_lens: torch.Tensor | None = None,
+        img_cu_lens: torch.Tensor | None = None,
+        # ms_pe: tuple[torch.FloatTensor, torch.FloatTensor] | None = None,
+        # pe: torch.FloatTensor | None = None,
+        # freqs_cos: torch.Tensor | None = None,
+        # freqs_sin: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        **attention_kwargs,
+    ) -> torch.Tensor:
+        r"""
+        The forward method of the `Attention` class.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                The hidden states of the query.
+            encoder_hidden_states (`torch.Tensor`, *optional*):
+                The hidden states of the encoder.
+            attention_mask (`torch.Tensor`, *optional*):
+                The attention mask to use. If `None`, no mask is applied.
+            **attention_kwargs:
+                Additional keyword arguments to pass along to the attention.
+
+        Returns:
+            `torch.Tensor`: The output of the attention layer.
+        """
+        # The `Attention` class can call different attention processors / attention functions
+        # here we simply pass along all tensors to the selected processor class
+        # For standard processors that are defined here, `**attention_kwargs` is empty
+
+        return self.processor(
+            self,
+            hidden_states,
+            attention_mask=attention_mask,
+            txt_cu_lens=txt_cu_lens,
+            img_cu_lens=img_cu_lens,
+            image_rotary_emb=image_rotary_emb,
+            **attention_kwargs,
+        )
+
+
+class MageDoubleStreamAttnProcessor:
+    """
+    Attention processor for the Mage double-stream architecture, matching DoubleStreamLayerMegatron logic. This processor
+    implements joint attention computation where text and image streams are processed together.
+    """
+
+    _attention_backend = None
+    _parallel_config = None
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "MageDoubleStreamAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,  # Image stream
+        img_cu_lens: torch.LongTensor,
+        attention_mask: torch.FloatTensor | None = None,
+        encoder_hidden_states: torch.FloatTensor = None,  # Text stream
+        txt_cu_lens: torch.LongTensor = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        if encoder_hidden_states is None:
+            raise ValueError("MageDoubleStreamAttnProcessor requires encoder_hidden_states (text stream)")
+
+        # seq_txt = encoder_hidden_states.shape[1]
+
+        # logger.info(f"hidden_states: {hidden_states.shape}")
+        # logger.info(f"encoder_hidden_states: {encoder_hidden_states.shape}")
+
+        # Compute QKV for image stream (sample projections). Experimental
+        # runtimes may provide one fused projection; the ordinary checkpoint
+        # path continues to use the three standard Linear modules.
+        if getattr(attn, "to_qkv", None) is not None:
+            img_query, img_key, img_value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
+        else:
+            img_query = attn.to_q(hidden_states)
+            img_key = attn.to_k(hidden_states)
+            img_value = attn.to_v(hidden_states)
+
+        # Compute QKV for text stream (context projections).
+        if getattr(attn, "add_qkv_proj", None) is not None:
+            txt_query, txt_key, txt_value = attn.add_qkv_proj(encoder_hidden_states).chunk(3, dim=-1)
+        else:
+            txt_query = attn.add_q_proj(encoder_hidden_states)
+            txt_key = attn.add_k_proj(encoder_hidden_states)
+            txt_value = attn.add_v_proj(encoder_hidden_states)
+
+        # Reshape for multi-head attention
+        img_query = img_query.unflatten(-1, (attn.heads, -1))
+        img_key = img_key.unflatten(-1, (attn.heads, -1))
+        img_value = img_value.unflatten(-1, (attn.heads, -1))
+
+        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
+        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
+        txt_value = txt_value.unflatten(-1, (attn.heads, -1))
+
+        # logger.info(
+        #     f"img_query shape: {img_query.shape}, img_key shape: {img_key.shape}, img_value shape: {img_value.shape}"
+        # )
+        # logger.info(
+        #     f"txt_query shape: {txt_query.shape}, txt_key shape: {txt_key.shape}, txt_value shape: {txt_value.shape}"
+        # )
+
+        if img_query.ndim == 4:
+            img_query = img_query.flatten(0, 1)
+            img_key = img_key.flatten(0, 1)
+            img_value = img_value.flatten(0, 1)
+
+        if txt_query.ndim == 4:
+            txt_query = txt_query.flatten(0, 1)
+            txt_key = txt_key.flatten(0, 1)
+            txt_value = txt_value.flatten(0, 1)
+
+        # Apply QK normalization
+        if attn.norm_q is not None:
+            img_query = attn.norm_q(img_query)
+        if attn.norm_k is not None:
+            img_key = attn.norm_k(img_key)
+        if attn.norm_added_q is not None:
+            txt_query = attn.norm_added_q(txt_query)
+        if attn.norm_added_k is not None:
+            txt_key = attn.norm_added_k(txt_key)
+
+        # logger.info(f"txt_query shape: {txt_query.shape}, txt_key shape: {txt_key.shape}")
+        # logger.info(f"freqs_cos shape: {freqs_cos.shape}, freqs_sin shape: {freqs_sin.shape}")
+
+        # Apply 2D multi-scale RoPE (MageFlowEmbedRope) to image tokens
+        img_freqs = image_rotary_emb
+        img_query = apply_rotary_emb_mageflow(img_query, img_freqs)
+        img_key = apply_rotary_emb_mageflow(img_key, img_freqs)
+        # Concatenate for joint attention
+        # Order: [text, image]
+        # joint_query = torch.cat([txt_query, img_query], dim=1)
+        # joint_key = torch.cat([txt_key, img_key], dim=1)
+        # joint_value = torch.cat([txt_value, img_value], dim=1)
+
+        # Calculate lengths
+        img_lens = img_cu_lens[1:] - img_cu_lens[:-1]
+        txt_lens = txt_cu_lens[1:] - txt_cu_lens[:-1]
+
+        # Calculate joint cu_seqlens
+        joint_lens = txt_lens + img_lens
+        joint_cu_lens = torch.cat(
+            [
+                torch.zeros(1, dtype=torch.int32, device=joint_lens.device),
+                torch.cumsum(joint_lens, dim=0, dtype=torch.int32),
+            ],
+            dim=0,
+        )
+
+        # logger.info(f"txt_lens: {txt_lens}, img_lens: {img_lens}")
+        # logger.info(f"joint_lens: {joint_lens}, joint_cu_lens: {joint_cu_lens}")
+
+        device = joint_lens.device
+        batch_size = len(txt_lens)
+        sample_indices = torch.arange(batch_size, device=device)
+
+        txt_sample_ids = torch.repeat_interleave(sample_indices, txt_lens)
+        img_sample_ids = torch.repeat_interleave(sample_indices, img_lens)
+
+        txt_intra_pos = torch.arange(txt_query.shape[0], device=device) - txt_cu_lens[txt_sample_ids]
+        img_intra_pos = torch.arange(img_query.shape[0], device=device) - img_cu_lens[img_sample_ids]
+
+        txt_dest_indices = joint_cu_lens[txt_sample_ids] + txt_intra_pos
+        img_dest_indices = joint_cu_lens[img_sample_ids] + txt_lens[img_sample_ids] + img_intra_pos
+
+        total_tokens = joint_cu_lens[-1]
+        joint_query = torch.empty((total_tokens, *txt_query.shape[1:]), dtype=txt_query.dtype, device=device)
+        joint_key = torch.empty((total_tokens, *txt_key.shape[1:]), dtype=txt_key.dtype, device=device)
+        joint_value = torch.empty((total_tokens, *txt_value.shape[1:]), dtype=txt_value.dtype, device=device)
+
+        # logger.info(f"joint_query shape: {joint_query.shape}")
+        # logger.info(f"joint_key shape: {joint_key.shape}")
+        # logger.info(f"joint_value shape: {joint_value.shape}")
+        # logger.info(f"txt_dest_indices shape: {txt_dest_indices.shape}")
+        # logger.info(f"img_dest_indices shape: {img_dest_indices.shape}")
+
+        joint_query[txt_dest_indices] = txt_query
+        joint_query[img_dest_indices] = img_query
+
+        joint_key[txt_dest_indices] = txt_key
+        joint_key[img_dest_indices] = img_key
+
+        joint_value[txt_dest_indices] = txt_value
+        joint_value[img_dest_indices] = img_value
+
+        max_seqlen = joint_lens.max().item()
+        joint_attn_output = flash_attn_varlen_func(
+            joint_query,
+            joint_key,
+            joint_value,
+            cu_seqlens_q=joint_cu_lens,
+            cu_seqlens_k=joint_cu_lens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            dropout_p=0.0,
+            softmax_scale=None,
+            causal=False,
+        )
+
+        txt_attn_output = joint_attn_output[txt_dest_indices]
+        img_attn_output = joint_attn_output[img_dest_indices]
+
+        img_attn_output = img_attn_output.flatten(1, 2)  # (N, H, D) -> (N, H*D)
+        img_attn_output = img_attn_output.to(joint_query.dtype)
+
+        txt_attn_output = txt_attn_output.flatten(1, 2)  # (N, H, D) -> (N, H*D)
+        txt_attn_output = txt_attn_output.to(joint_query.dtype)
+
+        img_attn_output = attn.to_out[0](img_attn_output)
+        if len(attn.to_out) > 1:
+            img_attn_output = attn.to_out[1](img_attn_output)  # dropout
+
+        txt_attn_output = attn.to_add_out(txt_attn_output)
+        txt_attn_output = txt_attn_output.view(
+            encoder_hidden_states.shape[0], encoder_hidden_states.shape[1], txt_attn_output.shape[-1]
+        )
+
+        return img_attn_output, txt_attn_output
+
+
+@maybe_allow_in_graph
+class MageFlowTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+
+        # Image processing modules
+        self.img_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,  # Enable cross attention for joint computation
+            added_kv_proj_dim=dim,  # Enable added KV projections for text stream
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=MageDoubleStreamAttnProcessor(),
+            eps=eps,
+        )
+        self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        # Text processing modules
+        self.txt_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        # Text doesn't need separate attention - it's handled by img_attn joint computation
+        self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+    def _modulate(self, x, mod_params, cu_lens=None, seq_lens=None):
+        """Apply modulation to input tensor"""
+        shift, scale, gate = mod_params.chunk(3, dim=-1)
+        if cu_lens is not None:
+            assert x.shape[0] == 1, "x must be of shape (1, *) when cu_lens is not None"
+            x_flattened = x.view(-1, x.shape[-1])
+            lengths = cu_lens[1:] - cu_lens[:-1]
+            shift_t = shift.repeat_interleave(lengths, dim=0)
+            scale_t = scale.repeat_interleave(lengths, dim=0)
+            gate_t = gate.repeat_interleave(lengths, dim=0)
+
+            x_flattened = x_flattened * (1 + scale_t) + shift_t
+            x = x_flattened.view(x.shape)
+            return x, gate_t
+        else:
+            return x * (1 + scale) + shift, gate
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        # encoder_hidden_states_mask: torch.Tensor,
+        temb: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        image_rotary_emb: torch.Tensor,
+        # freqs_cos: torch.Tensor,
+        # freqs_sin: torch.Tensor,
+        txt_cu_lens: torch.Tensor,
+        img_cu_lens: torch.Tensor,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Get modulation parameters for both streams
+        # if isinstance(temb, tuple):
+        #     temb_img, temb_txt = temb
+        # else:
+        #     temb_img = temb_txt = temb
+
+        img_mod_params = self.img_mod(temb)  # [B, 6*dim]
+        txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
+
+        # logger.info(f"img_mod_params: {img_mod_params.shape}, txt_mod_params: {txt_mod_params.shape}")
+
+        # if img_cu_lens is not None and txt_cu_lens is not None and hidden_states.ndim == 2:
+        #     img_lens = img_cu_lens[1:] - img_cu_lens[:-1]
+        #     txt_lens = txt_cu_lens[1:] - txt_cu_lens[:-1]
+        #     img_mod_params = img_mod_params.repeat_interleave(img_lens, dim=0)
+        #     txt_mod_params = txt_mod_params.repeat_interleave(txt_lens, dim=0)
+
+        # Split modulation parameters for norm1 and norm2
+        img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+        txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+
+        # Process image stream - norm1 + modulation
+        img_normed = self.img_norm1(hidden_states)
+        img_modulated, img_gate1 = self._modulate(img_normed, img_mod1, cu_lens=img_cu_lens)
+
+        # Process text stream - norm1 + modulation
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1, cu_lens=txt_cu_lens)
+
+        # Use MageDoubleStreamAttnProcessor for joint attention computation
+        # This directly implements the DoubleStreamLayerMegatron logic:
+        # 1. Computes QKV for both streams
+        # 2. Applies QK normalization and RoPE
+        # 3. Concatenates and runs joint attention
+        # 4. Splits results back to separate streams
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        # logger.info(f"img_modulated: {img_modulated}")
+        # logger.info(f"txt_modulated: {txt_modulated}")
+        attn_output = self.attn(
+            hidden_states=img_modulated,  # Image stream (will be processed as "sample")
+            encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
+            # encoder_hidden_states_mask=encoder_hidden_states_mask,
+            image_rotary_emb=image_rotary_emb,
+            txt_cu_lens=txt_cu_lens,
+            img_cu_lens=img_cu_lens,
+            # freqs_cos=freqs_cos,
+            # freqs_sin=freqs_sin,
+            **joint_attention_kwargs,
+        )
+        # logger.info(f"attn_output: {attn_output}")
+
+        # MageDoubleStreamAttnProcessor returns (img_output, txt_output) when encoder_hidden_states is provided
+        img_attn_output, txt_attn_output = attn_output
+
+        # Apply attention gates and add residual (like in Megatron)
+        hidden_states = hidden_states + img_gate1 * img_attn_output
+        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
+
+        # Process image stream - norm2 + MLP
+        img_normed2 = self.img_norm2(hidden_states)
+        img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2, cu_lens=img_cu_lens)
+        img_mlp_output = self.img_mlp(img_modulated2)
+        hidden_states = hidden_states + img_gate2 * img_mlp_output
+
+        # Process text stream - norm2 + MLP
+        txt_normed2 = self.txt_norm2(encoder_hidden_states)
+        txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2, cu_lens=txt_cu_lens)
+        txt_mlp_output = self.txt_mlp(txt_modulated2)
+        encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
+
+        # Clip to prevent overflow for fp16
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class AdaLayerNormContinuous(nn.Module):
+    r"""
+    Adaptive normalization layer with a norm layer (layer_norm or rms_norm).
+
+    Args:
+        embedding_dim (`int`): Embedding dimension to use during projection.
+        conditioning_embedding_dim (`int`): Dimension of the input condition.
+        elementwise_affine (`bool`, defaults to `True`):
+            Boolean flag to denote if affine transformation should be applied.
+        eps (`float`, defaults to 1e-5): Epsilon factor.
+        bias (`bias`, defaults to `True`): Boolean flag to denote if bias should be use.
+        norm_type (`str`, defaults to `"layer_norm"`):
+            Normalization layer to use. Values supported: "layer_norm", "rms_norm".
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
+        # because the output is immediately scaled and shifted by the projected conditioning embeddings.
+        # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
+        # However, this is how it was implemented in the original code, and it's rather likely you should
+        # set `elementwise_affine` to False.
+        elementwise_affine=True,
+        eps=1e-5,
+        bias=True,
+        norm_type="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":
+            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(embedding_dim, eps, elementwise_affine)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(
+        self, x: torch.Tensor, conditioning_embedding: torch.Tensor,
+        cu_seqlens: torch.Tensor | None = None, seq_lens: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for
+        # hunyuanDiT)
+        emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
+        if cu_seqlens is None:
+            scale, shift = torch.chunk(emb, 2, dim=-1)
+            x = self.norm(x) * (1 + scale) + shift
+        else:
+            sample_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            flattened_x = x.view(-1, x.shape[-1])
+            scale, shift = torch.chunk(emb, 2, dim=-1)
+            scale_t = torch.repeat_interleave(scale, sample_lens, dim=0)
+            shift_t = torch.repeat_interleave(shift, sample_lens, dim=0)
+            flattened_x = self.norm(flattened_x) * (1 + scale_t) + shift_t
+            x = flattened_x.view(x.shape)
+        return x