Add vendor/mage_flow/models/mage_flow.py

vendor/mage_flow/models/mage_flow.py

@@ -0,0 +1,364 @@
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from loguru import logger
+from pydantic import BaseModel, Field
+from torch import Tensor
+
+from .modules._attn_backend import set_attn_backend
+from .modules.mage_layers import (
+    AdaLayerNormContinuous,
+    MageFlowEmbedRope,
+    MageFlowTimestepProjEmbeddings,
+    MageFlowTransformerBlock,
+    RMSNorm,
+)
+from .modules.text_encoder import TextEncoder, qwen3_patch_forward
+
+
+class ModelConfig(BaseModel):
+    static_shift: float = Field(
+        default=6.0,
+        description="Static shift value for the z-image time-shift schedule (the only "
+                    "supported schedule). Default: 6.0.",
+    )
+    vae_path: str = Field(...)
+    model_structure: dict = Field(default_factory=dict)
+    txt_enc_path: str = Field(...)
+    txt_max_length: int = Field(default=4096)
+    pretrained_model_name_or_path: str | None = Field(default=None)
+    pretrained_full_model_path: str | None = Field(default=None)  # Load full model weights (DiT + txt_enc + vae)
+    packing: bool = Field(default=False)
+    vae_sample_posterior: bool = Field(default=True)   # Sample (vs mode) from VAE posterior at encode time (Flux2 + CoD)
+    vae_encoder_only: bool = Field(default=False)  # Skip loading VAE decoder to save GPU memory (training only, MageVAE)
+    compile_vae_encoder: bool = Field(default=False)  # torch.compile VAE encoder to reduce CUDA kernel launch overhead
+    attn_type: str = Field(
+        default="flash2",
+        description="Flash-attn backend used by both the DiT (mage_layers) "
+                    "and the HF text encoder (text_encoder). One of: 'flash2' (default) or 'flash4'.",
+    )
+
+
+@dataclass
+class MageFlowParams:
+    in_channels: int
+    out_channels: int
+    context_in_dim: int
+    hidden_size: int
+    num_heads: int
+    depth: int
+    axes_dim: list[int]
+    checkpoint: bool
+    patch_size: int = 1
+
+
+class MageFlow(nn.Module):
+    def __init__(self, params: MageFlowParams):
+        super().__init__()
+        self.params = params
+        self.checkpoint = params.checkpoint
+        self.in_channels = params.in_channels
+        self.out_channels = params.out_channels
+        self.inner_dim = params.hidden_size  # num_attention_heads * attention_head_dim
+        self.axes_dim = params.axes_dim
+        self.num_attention_heads = params.num_heads
+        self.attention_head_dim = self.inner_dim // self.num_attention_heads
+        self.patch_size = params.patch_size
+        assert sum(self.axes_dim) == self.attention_head_dim
+
+        self.pos_embed = MageFlowEmbedRope(theta=10000, axes_dim=self.axes_dim, scale_rope=True)
+        self.img_in = nn.Linear(self.in_channels, self.inner_dim)
+        self.txt_norm = RMSNorm(params.context_in_dim, eps=1e-6)
+        self.txt_in = nn.Linear(params.context_in_dim, self.inner_dim)
+
+        self.time_text_embed = MageFlowTimestepProjEmbeddings(embedding_dim=self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MageFlowTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, self.patch_size * self.patch_size * self.out_channels, bias=True)
+
+    def forward(
+        self,
+        img: Tensor,
+        txt: Tensor,
+        timesteps: Tensor,
+        img_shapes=None,
+        img_cu_seqlens: Tensor | None = None,
+        txt_cu_seqlens: Tensor | None = None,
+        attention_kwargs: dict[str, Any] | None = None,
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # Prepare vision RoPE (msrope); text tokens are not rotated.
+        ms_pe = self.pos_embed(img_shapes, device=img.device)
+
+        img = self.img_in(img)
+        txt = self.txt_norm(txt)
+
+        timesteps = timesteps.to(img.dtype)
+        temb = self.time_text_embed(timesteps, img)
+
+        txt = self.txt_in(txt)
+        txt_vec = torch.zeros(txt.shape[0], self.inner_dim, dtype=txt.dtype, device=txt.device)
+
+        temb = temb + txt_vec
+
+        attention_kwargs = attention_kwargs or {}
+
+        for _index_block, block in enumerate(self.transformer_blocks):
+            if self.training and self.checkpoint:
+                txt, img = torch.utils.checkpoint.checkpoint(
+                    block,
+                    img,  # hidden_states
+                    txt,  # encoder_hidden_states
+                    temb,  # temb
+                    ms_pe,  # image_rotary_emb
+                    txt_cu_seqlens,  # txt_cu_lens
+                    img_cu_seqlens,  # img_cu_lens
+                    use_reentrant=False,
+                )
+
+            else:
+                txt, img = block(
+                    hidden_states=img,
+                    encoder_hidden_states=txt,
+                    txt_cu_lens=txt_cu_seqlens,
+                    img_cu_lens=img_cu_seqlens,
+                    temb=temb,
+                    image_rotary_emb=ms_pe,
+                    joint_attention_kwargs=attention_kwargs,
+                )
+
+        # Use only the image part (hidden_states) from the dual-stream blocks
+        img = self.norm_out(
+            img,
+            temb,
+            cu_seqlens=img_cu_seqlens,
+        )
+        img = self.proj_out(img)
+        return img
+
+
+class MageFlowModel(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        set_attn_backend(getattr(config, "attn_type", "flash2"))
+        self.patch_text_encoder_forward()
+        self.vae = self.load_vae()
+        self.transformer = self.load_transformer()
+        self.txt_enc = self.load_text_enc()
+
+        # Optionally override all components from a full model checkpoint (e.g. ema.pt)
+        full_path = getattr(self.config, "pretrained_full_model_path", None)
+        if full_path is not None:
+            import os
+
+            if os.path.exists(full_path):
+                logger.info(f"Loading full model weights from {full_path}")
+                sd = torch.load(full_path, map_location="cpu")
+                # Handle wrapped EMA format: {'ema_state_dict': ..., ...}
+                if isinstance(sd, dict) and "ema_state_dict" in sd:
+                    sd = sd["ema_state_dict"]
+                missing, unexpected = self.load_state_dict(sd, strict=False)
+                if missing:
+                    logger.warning(f"Full model load missing keys ({len(missing)}): {missing[:5]}...")
+                if unexpected:
+                    logger.warning(f"Full model load unexpected keys ({len(unexpected)}): {unexpected[:5]}...")
+                logger.info("Full model weights loaded successfully.")
+            else:
+                logger.warning(f"pretrained_full_model_path not found: {full_path}")
+
+        # Freeze VAE and Text Encoder
+        self.vae.requires_grad_(False)
+
+        # Drop VAE decoder to save GPU memory (training only, decoder unused during training)
+        if self.config.vae_encoder_only:
+            from .modules.mage_vae import MageVAE
+            if isinstance(self.vae, MageVAE):
+                decoder_params = sum(p.numel() for p in self.vae.decoder_model.parameters()) / 1e6
+                self.vae.decoder_model = None
+            elif hasattr(self.vae, "decoder"):
+                decoder_params = sum(p.numel() for p in self.vae.decoder.parameters()) / 1e6
+                self.vae.decoder = None
+            else:
+                decoder_params = 0
+            logger.info(f"vae_encoder_only=True: dropped VAE decoder ({decoder_params:.1f}M params) to save memory")
+
+        # NOTE: VAE encoder torch.compile() is deferred to
+        # maybe_compile_vae_encoder(), called after checkpoint load. Reason:
+        # avoid wasted compile work before load_checkpoint overwrites weights.
+        # The save-side _unwrap_compiled_submodules guard in DeepSpeedTrainer
+        # is a belt-and-suspenders defense against any future code that
+        # re-introduces the function-style ``module = torch.compile(module)``
+        # pattern (which does pollute state_dict with ``_orig_mod.``).
+
+        # Text encoder is always frozen (inference only).
+        self.txt_enc.requires_grad_(False)
+        logger.info(
+            f"{sum([p.numel() for p in self.transformer.parameters() if p.requires_grad]) / 1000000} M parameters"
+        )
+
+    def patch_text_encoder_forward(self):
+        qwen3_patch_forward()
+        logger.info("Patched Qwen3-VL text encoder forward methods")
+
+    def maybe_compile_vae_encoder(self) -> None:
+        """Compile the VAE encoder with torch.compile() to fuse small ops and
+        reduce CUDA kernel launch overhead.
+
+        Uses ``nn.Module.compile()`` (in-place) for the encoder so the module
+        hierarchy and parameter names are unchanged — ``state_dict()`` keeps
+        clean keys (no ``_orig_mod.`` prefix), and checkpoints stay
+        interchangeable with the non-compiled path.
+
+        For the MageVAE branch we still assign ``torch.compile(...)`` to a
+        method (``_encode_moments``); methods aren't ``nn.Module``s so this
+        does not pollute ``state_dict()``.
+
+        Idempotent: safe to call multiple times; already-compiled modules are
+        detected and skipped.
+        """
+        if not getattr(self.config, "compile_vae_encoder", False):
+            return
+        torch.set_float32_matmul_precision("high")
+        from .modules.mage_vae import MageVAE
+        if isinstance(self.vae, MageVAE):
+            fn = self.vae._encode_moments
+            if hasattr(fn, "_torchdynamo_orig_callable") or hasattr(fn, "_orig_mod"):
+                return  # already compiled
+            self.vae._encode_moments = torch.compile(fn, dynamic=True)
+            logger.info("compile_vae_encoder=True: compiled MageVAE._encode_moments")
+        elif hasattr(self.vae, "encoder"):
+            if getattr(self.vae.encoder, "_compiled_call_impl", None) is not None:
+                return  # already compiled
+            self.vae.encoder.compile()
+            logger.info("compile_vae_encoder=True: compiled VAE encoder (in-place)")
+
+    def load_text_enc(self):
+        return TextEncoder(
+            model_name=self.config.txt_enc_path,
+            version=self.config.txt_enc_path,
+            tokenizer_max_length=self.config.txt_max_length,
+            torch_dtype=torch.bfloat16,
+            prompt_template=None,
+            dit_structure=self.config.model_structure,
+            use_packed_text_infer=self.config.packing,
+            attn_type=getattr(self.config, "attn_type", "flash2"),
+        )
+
+    def load_vae(self):
+        from .modules.mage_vae import MageVAE
+        return MageVAE(
+            ckpt_path=self.config.vae_path,
+            sample_posterior=self.config.vae_sample_posterior,
+        )
+
+    def load_transformer(self):
+        # Imported lazily to avoid a circular import: ``utils`` imports MageFlow /
+        # MageFlowParams from this module.
+        from .utils import load_model
+        return load_model(
+            dit_structure=self.config.model_structure,
+            pretrain_path=self.config.pretrained_model_name_or_path,
+        )
+
+    def compile(self):
+        self.transformer.compile()
+
+    def compute_vae_encodings(
+        self,
+        pixel_values: torch.Tensor | list[torch.Tensor],
+        with_ids: bool = True,
+    ):
+        if isinstance(pixel_values, list):
+            # All same resolution → batch encode via the tensor path
+            if len(pixel_values) > 1 and len({img.shape for img in pixel_values}) == 1:
+                stacked = torch.stack(pixel_values, dim=0)
+                result = self.compute_vae_encodings(stacked, with_ids=with_ids)
+                # Repack from [N, L, C] batch format to [1, N*L, C] packed format
+                if with_ids:
+                    model_input, img_shapes, img_ids = result
+                    model_input = model_input.reshape(1, -1, model_input.shape[-1])
+                    img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
+                    return model_input, img_shapes, img_ids
+                model_input, img_shapes = result
+                model_input = model_input.reshape(1, -1, model_input.shape[-1])
+                return model_input, img_shapes
+
+            # Packed / variable-size images
+            model_inputs = []
+            img_shapes = []
+            img_ids_list = []
+
+            def _append(latents):
+                _, _, h, w = latents.shape
+                img_shapes.append([(1, h, w)])
+                model_inputs.append(rearrange(latents, "b c h w -> b (h w) c").squeeze(0))
+                if with_ids:
+                    ids = torch.zeros(h, w, 3, device=latents.device)
+                    ids[..., 1] = ids[..., 1] + torch.arange(h, device=latents.device)[:, None]
+                    ids[..., 2] = ids[..., 2] + torch.arange(w, device=latents.device)[None, :]
+                    img_ids_list.append(rearrange(ids, "h w c -> (h w) c"))
+
+            # MageVAE encoder is launch-bound on B=1; group same-shape images
+            # in the pack into one batched encode call.
+            if len(pixel_values) > 1:
+                groups: dict[tuple[int, int], list[int]] = {}
+                for i, img in enumerate(pixel_values):
+                    key = (int(img.shape[-2]), int(img.shape[-1]))
+                    groups.setdefault(key, []).append(i)
+                latents_per_idx = [None] * len(pixel_values)
+                for (h, w), idxs in groups.items():
+                    batch = torch.stack([pixel_values[i] for i in idxs], dim=0)
+                    batch = batch.to(memory_format=torch.contiguous_format).float()
+                    batch = batch.to(self.vae.device, dtype=self.vae.dtype)
+                    with torch.no_grad():
+                        lat = self.vae.encode(batch)  # [B, 128, H/16, W/16]
+                    for j, i in enumerate(idxs):
+                        latents_per_idx[i] = lat[j:j + 1]
+                for latents in latents_per_idx:
+                    _append(latents)
+            else:
+                for img in pixel_values:
+                    img = img.unsqueeze(0).to(memory_format=torch.contiguous_format).float()
+                    img = img.to(self.vae.device, dtype=self.vae.dtype)
+                    with torch.no_grad():
+                        latents = self.vae.encode(img)  # [1, 128, H/16, W/16]
+                    _append(latents)
+
+            model_input = torch.cat(model_inputs, dim=0).unsqueeze(0)
+            if with_ids:
+                img_ids = torch.cat(img_ids_list, dim=0).unsqueeze(0)
+                return model_input, img_shapes, img_ids
+            return model_input, img_shapes
+
+        # Tensor (padded batch)
+        pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
+        pixel_values = pixel_values.to(self.vae.device, dtype=self.vae.dtype)
+        with torch.no_grad():
+            model_input = self.vae.encode(pixel_values)  # [B, 128, H/16, W/16]
+        bs, c, h, w = model_input.shape
+        img_shapes = [[(1, h, w)]] * bs
+        model_input = rearrange(model_input, "b c h w -> b (h w) c")
+        if with_ids:
+            img_ids = torch.zeros(h, w, 3, device=model_input.device)
+            img_ids[..., 1] = img_ids[..., 1] + torch.arange(h, device=model_input.device)[:, None]
+            img_ids[..., 2] = img_ids[..., 2] + torch.arange(w, device=model_input.device)[None, :]
+            img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+            return model_input, img_shapes, img_ids
+        return model_input, img_shapes