ComfyUI/comfy/ldm/mage_flow/model.py

# Mage-Flow (https://github.com/microsoft/Mage) native-resolution MMDiT (MIT)
# Architecture is a 12-layer variant of the Qwen-Image double-stream block with
# patch_size=1 (no 2x2 packing), unrotated text tokens and a bf16-rounded
# timestep frequency table.
import math
import torch
import torch.nn as nn
from typing import Optional, Tuple

from comfy.ldm.lightricks.model import TimestepEmbedding
from comfy.ldm.flux.layers import EmbedND
from comfy.ldm.qwen_image.model import QwenImageTransformerBlock, LastLayer
import comfy.patcher_extension


class MageTimestepProjEmbeddings(nn.Module):
    def __init__(self, embedding_dim, dtype=None, device=None, operations=None):
        super().__init__()
        self.timestep_embedder = TimestepEmbedding(
            in_channels=256, time_embed_dim=embedding_dim,
            dtype=dtype, device=device, operations=operations
        )

    def forward(self, timestep, hidden_states):
        half_dim = 128
        exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timestep.device) / half_dim
        emb = torch.exp(exponent).to(timestep.dtype)
        emb = timestep[:, None].float() * emb[None, :]
        emb = 1000.0 * emb
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)  # flip_sin_to_cos
        return self.timestep_embedder(emb.to(dtype=hidden_states.dtype))


class MageFlowTransformer2DModel(nn.Module):
    def __init__(
        self,
        in_channels: int = 128,
        out_channels: Optional[int] = 128,
        num_layers: int = 12,
        attention_head_dim: int = 128,
        num_attention_heads: int = 24,
        joint_attention_dim: int = 2560,
        axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
        image_model=None,
        dtype=None,
        device=None,
        operations=None,
    ):
        super().__init__()
        self.dtype = dtype
        self.patch_size = 1
        self.in_channels = in_channels
        self.out_channels = out_channels or in_channels
        self.inner_dim = num_attention_heads * attention_head_dim

        self.pe_embedder = EmbedND(dim=attention_head_dim, theta=10000, axes_dim=list(axes_dims_rope))

        self.time_text_embed = MageTimestepProjEmbeddings(embedding_dim=self.inner_dim, dtype=dtype, device=device, operations=operations)

        self.txt_norm = operations.RMSNorm(joint_attention_dim, eps=1e-6, dtype=dtype, device=device)
        self.img_in = operations.Linear(in_channels, self.inner_dim, dtype=dtype, device=device)
        self.txt_in = operations.Linear(joint_attention_dim, self.inner_dim, dtype=dtype, device=device)

        self.transformer_blocks = nn.ModuleList([
            QwenImageTransformerBlock(
                dim=self.inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                dtype=dtype,
                device=device,
                operations=operations
            )
            for _ in range(num_layers)
        ])

        self.norm_out = LastLayer(self.inner_dim, self.inner_dim, dtype=dtype, device=device, operations=operations)
        self.proj_out = operations.Linear(self.inner_dim, self.out_channels, bias=True, dtype=dtype, device=device)

    def process_img(self, x, index=0):
        # patch_size=1: tokens are raw latent pixels, no 2x2 packing.
        bs, c, h, w = x.shape
        hidden_states = x.movedim(1, -1).reshape(bs, h * w, c)

        img_ids = torch.zeros((h, w, 3), device=x.device)
        # Frame axis: positive image index (0 = target, 1..N = reference images).
        img_ids[:, :, 0] = index
        # Mage scale_rope centering: positions [-ceil(n/2), floor(n/2)), i.e.
        # offset by (n - n//2). Differs from Qwen-Image's -(n//2) for odd sizes.
        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.arange(h, device=x.device)[:, None] - (h - h // 2)
        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.arange(w, device=x.device)[None, :] - (w - w // 2)
        return hidden_states, img_ids.reshape(h * w, 3).unsqueeze(0).expand(bs, -1, -1), (h, w)

    def forward(self, x, timestep, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
        return comfy.patcher_extension.WrapperExecutor.new_class_executor(
            self._forward,
            self,
            comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
        ).execute(x, timestep, context, attention_mask, ref_latents, transformer_options, **kwargs)

    def _forward(self, x, timestep, context, attention_mask=None, ref_latents=None, transformer_options={}, control=None, **kwargs):
        if attention_mask is not None and not torch.is_floating_point(attention_mask):
            attention_mask = (attention_mask - 1).to(x.dtype) * torch.finfo(x.dtype).max

        hidden_states, img_ids, orig_shape = self.process_img(x)
        num_embeds = hidden_states.shape[1]

        if ref_latents is not None:
            ref_num_tokens = []
            index = 0
            for ref in ref_latents:
                index += 1
                kontext, kontext_ids, _ = self.process_img(ref, index=index)
                hidden_states = torch.cat([hidden_states, kontext], dim=1)
                img_ids = torch.cat([img_ids, kontext_ids], dim=1)
                ref_num_tokens.append(kontext.shape[1])
            transformer_options = transformer_options.copy()
            transformer_options["reference_image_num_tokens"] = ref_num_tokens

        # Text tokens are not rotated in Mage-Flow: RoPE at position 0 is the
        # identity rotation.
        txt_ids = torch.zeros((x.shape[0], context.shape[1], 3), device=x.device)

        hidden_states = self.img_in(hidden_states)
        context = self.txt_norm(context)
        context = self.txt_in(context)

        temb = self.time_text_embed(timestep, hidden_states)

        patches_replace = transformer_options.get("patches_replace", {})
        patches = transformer_options.get("patches", {})
        blocks_replace = patches_replace.get("dit", {})

        if "post_input" in patches:
            for p in patches["post_input"]:
                out = p({"img": hidden_states, "txt": context, "img_ids": img_ids, "txt_ids": txt_ids, "transformer_options": transformer_options})
                hidden_states = out["img"]
                context = out["txt"]
                img_ids = out["img_ids"]
                txt_ids = out["txt_ids"]

        ids = torch.cat((txt_ids, img_ids), dim=1)
        image_rotary_emb = self.pe_embedder(ids).contiguous()
        del ids, txt_ids, img_ids

        transformer_options["total_blocks"] = len(self.transformer_blocks)
        transformer_options["block_type"] = "double"
        for i, block in enumerate(self.transformer_blocks):
            transformer_options["block_index"] = i
            if ("double_block", i) in blocks_replace:
                def block_wrap(args):
                    out = {}
                    out["txt"], out["img"] = block(hidden_states=args["img"], encoder_hidden_states=args["txt"], encoder_hidden_states_mask=attention_mask, temb=args["vec"], image_rotary_emb=args["pe"], transformer_options=args["transformer_options"])
                    return out
                out = blocks_replace[("double_block", i)]({"img": hidden_states, "txt": context, "vec": temb, "pe": image_rotary_emb, "transformer_options": transformer_options}, {"original_block": block_wrap})
                hidden_states = out["img"]
                context = out["txt"]
            else:
                context, hidden_states = block(
                    hidden_states=hidden_states,
                    encoder_hidden_states=context,
                    encoder_hidden_states_mask=attention_mask,
                    temb=temb,
                    image_rotary_emb=image_rotary_emb,
                    transformer_options=transformer_options,
                )

            if "double_block" in patches:
                for p in patches["double_block"]:
                    out = p({"img": hidden_states, "txt": context, "x": x, "block_index": i, "transformer_options": transformer_options})
                    hidden_states = out["img"]
                    context = out["txt"]

            if control is not None:  # Controlnet
                control_i = control.get("input")
                if i < len(control_i):
                    add = control_i[i]
                    if add is not None:
                        hidden_states[:, :add.shape[1]] += add

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

        hidden_states = hidden_states[:, :num_embeds]
        h, w = orig_shape
        return hidden_states.reshape(x.shape[0], h, w, self.out_channels).movedim(-1, 1)