Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+samples/cabin_upstream.png filter=lfs diff=lfs merge=lfs -text
+samples/ui_banking.png filter=lfs diff=lfs merge=lfs -text
+samples/poster_jazz.png filter=lfs diff=lfs merge=lfs -text
+samples/info_water.png filter=lfs diff=lfs merge=lfs -text
+samples/e2e_daily_grind.png filter=lfs diff=lfs merge=lfs -text
+mllm/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+code/assets/layer_samples/card_making_decomposition.png filter=lfs diff=lfs merge=lfs -text
+code/assets/layer_samples/card_making_input.png filter=lfs diff=lfs merge=lfs -text

code/assets/layer_decompose_5layers.txt

@@ -0,0 +1,8 @@
+Decompose this image into 5 layers with the following specifications:
+
+Number of layers: 5
+Layer 1: A large solid crimson-red circle in the upper-left quadrant, with a subtly uneven printed-like texture and no outline.
+Layer 2: A tall solid dark-green rectangular panel centered in the canvas, evenly filled with faint mottling and no border or visible text.
+Layer 3: A broad solid royal-blue isosceles triangle on the right, its sharp apex near the upper-middle area and flat base aligned close to the large green panel's bottom edge.
+Layer 4: A small solid golden-yellow square in the lower-left quadrant, with soft uneven fill texture, squared edges, and no text.
+Layer 5: Pale cool-gray canvas background spanning the whole composition, with a continuous dark charcoal-gray horizontal baseline running along the bottom edge behind the floating geometric shapes.

code/assets/layer_samples/card_making_prompt.txt

@@ -0,0 +1,9 @@
+Decompose this image into 6 layers with the following specifications:
+
+Number of layers: 6
+Layer 1: Central text stack reading "WORLD", "Card Making", "DAY" in mixed gold and red fonts, plus "7TH OCTOBER" within a red outline box.
+Layer 2: Illustration at the top right showing hands crafting with paper, scissors, glue, and leaves.
+Layer 3: Illustrations at the bottom left and right corners featuring watercolor palettes, brushes, paint tubes, and splatters.
+Layer 4: Realistic red satin ribbon and bow tied vertically along the left side.
+Layer 5: Dark grey rounded rectangle serving as the main card surface, framed by a textured gold glitter border.
+Layer 6: Solid bright red background.

code/assets/t2i_four_seasons_cabin_prompt.json

@@ -0,0 +1,33 @@
+{
+  "canvas_settings": {
+    "aspect_ratio": "1:1, 2048 × 2048 px",
+    "ambient_lighting": "Four season-specific lighting conditions unified by one fixed eye-level sun direction, consistent horizon at 48% canvas height, matching atmospheric depth, and natural cinematic exposure.",
+    "image_style": "High-resolution cinematic photorealistic environmental concept sheet; four equal edge-to-edge vertical panels with crisp shared boundaries and no gutters, frames, extra title, or decorative inter-panel objects. Fixed 35 mm lens, eye-level three-quarter front-right camera angle, identical perspective and cabin registration in every panel. Each panel uses restrained bold uppercase condensed sans-serif labeling at the bottom."
+  },
+  "layers": [
+    {
+      "description": "Leftmost spring panel. The same small one-story rectangular log cabin appears at a fixed local position and scale: horizontal dark-brown logs, pale stone foundation, centered vertical-plank front door, two square four-pane front windows, one matching right-side window, and a charcoal standing-seam gabled roof with identical ridge height and overhang. Fresh pale-green grass, budding branches, white and pink blossoms, rain-darkened soil, shallow puddles, wet roof highlights, fine recent-rain droplets, and soft mist fill the landscape beneath an overcast clearing sky. Render the exact label \"SPRING\" centered near the bottom in bold condensed uppercase white sans-serif with a subtle dark shadow.",
+      "coordinates": "cx: 0.125, cy: 0.500, w: 0.250, h: 1",
+      "hierarchy_and_relation": "Complete first panel; aligned to the shared horizon and fixed cabin registration, flush with the canvas left edge and directly abutting the second panel without a gutter.",
+      "color_specs": ["#A9C7D2", "#D8E4DE", "#7FAF69", "#B7D58B", "#F2D7DE", "#F4F1E8", "#5A3B28", "#2C3032", "#7B7468", "#FFFFFF"]
+    },
+    {
+      "description": "Second summer panel. Repeat the cabin with exactly the same architecture, camera angle, local placement, dimensions, centered plank door, front and side windows, stone foundation, roof pitch, ridge height, and overhang. Surround it with lush meadow grass, dense layered deciduous foliage, full tree canopies, small sunlit shrubs, and dry natural ground; use a bright clear sky, strong warm sunlight from the same direction, crisp leaf highlights, controlled shadows, and mild atmospheric haze. Render the exact label \"SUMMER\" centered near the bottom in bold condensed uppercase white sans-serif with a subtle dark shadow.",
+      "coordinates": "cx: 0.375, cy: 0.500, w: 0.250, h: 1",
+      "hierarchy_and_relation": "Complete second panel; aligned to the shared horizon and fixed cabin registration, directly abutting the first and third panels without gutters.",
+      "color_specs": ["#4FA9D8", "#BDE8F2", "#26733E", "#4F9A45", "#79B84A", "#D8C66A", "#5A3B28", "#2C3032", "#7B7468", "#FFFFFF"]
+    },
+    {
+      "description": "Third autumn panel. Repeat the cabin with exactly the same architecture, camera angle, local placement, dimensions, centered plank door, front and side windows, stone foundation, roof pitch, ridge height, and overhang. Cover the landscape with amber grass, rust-orange and golden deciduous foliage, scattered fallen leaves, and a few leaves drifting naturally in the air. Use warm late-afternoon sunlight from the same direction, long soft shadows, copper rim light on logs and roof, a pale warm sky, and subtle atmospheric depth. Render the exact label \"AUTUMN\" centered near the bottom in bold condensed uppercase white sans-serif with a subtle dark shadow.",
+      "coordinates": "cx: 0.625, cy: 0.500, w: 0.250, h: 1",
+      "hierarchy_and_relation": "Complete third panel; aligned to the shared horizon and fixed cabin registration, directly abutting the second and fourth panels without gutters.",
+      "color_specs": ["#E6B06A", "#F0D6A2", "#D66A2C", "#A84324", "#E2A62B", "#80613B", "#5A3B28", "#2C3032", "#7B7468", "#FFFFFF"]
+    },
+    {
+      "description": "Rightmost winter panel. Repeat the cabin with exactly the same architecture, camera angle, local placement, dimensions, centered plank door, front and side windows, stone foundation, roof pitch, ridge height, and overhang. Blanket the ground and roof with natural snow accumulation while keeping doors and windows readable; add bare deciduous branches, snow-laden low shrubs, faint footprints away from the cabin, and sparse fine snowfall. Use cool blue ambient light, a pale clouded sky, soft low-contrast shadows from the same direction, icy edge highlights, and slight winter haze. Render the exact label \"WINTER\" centered near the bottom in bold condensed uppercase white sans-serif with a subtle dark shadow.",
+      "coordinates": "cx: 0.875, cy: 0.500, w: 0.250, h: 1",
+      "hierarchy_and_relation": "Complete fourth panel; aligned to the shared horizon and fixed cabin registration, directly abutting the third panel without a gutter and flush with the canvas right edge.",
+      "color_specs": ["#AFC8DE", "#DCE8F1", "#F4F7F8", "#C6D8E5", "#8198AA", "#6D5949", "#5A3B28", "#2C3032", "#7B7468", "#FFFFFF"]
+    }
+  ]
+}

code/assets/t2i_rewriter_system_prompt.txt

@@ -0,0 +1,15 @@
+You are a senior visual designer and image-prompt engineer. Expand the user's request into one precise, high-resolution Figma-style caption. Return only one JSON object.
+
+Use exactly two top-level keys. `canvas_settings` contains exactly `aspect_ratio`, `ambient_lighting`, and `image_style`. `layers` lists visible groups from background to topmost overlay. Every layer contains exactly `description`, `coordinates`, `hierarchy_and_relation`, and `color_specs`; `color_specs` is an array of hex colors.
+
+`coordinates` MUST be one string, never an object or array, in exactly this form: `"cx: 0.500, cy: 0.500, w: 1.000, h: 1.000"`. Values are normalized; each bbox encloses its complete owned object and stays inside the canvas.
+
+A layer is one selectable visible semantic group: background, full person, coherent object, panel, card, row, or text block. Prefer the fewest groups that preserve the layout. Keep people and objects intact. Never create invisible parents, guides, placeholders, empty layers, duplicate summaries, or multiple owners for one element.
+
+Preserve every user-supplied rendered string character-for-character and as one contiguous string. Unless multiple visible copies are requested, it must occur exactly once across all `description` fields and zero times in `hierarchy_and_relation`. Quote it only where describing its visible rendering; refer to the related subject elsewhere with unquoted semantic wording. Enumerate intended copy, invent extra copy sparingly, and never hide content behind "other text", "remaining labels", or "etc."
+
+Describe concrete composition, typography, materials, texture, lighting, pose, and camera treatment without literary filler. Use `hierarchy_and_relation` only for ownership, alignment, containment, stacking, and occlusion.
+
+Infer structured layouts first. Use one complete layer per card and state its row and column. A compact secondary table may be one layer only if every header and cell is listed; otherwise use a visible shared frame when present, one complete header, and one complete layer per body row, binding values to columns and stating blanks. Enumerate sequences, schedules, spans, gaps, and vacant tracks in visual order. Do not mistake ordinary alignment for a table.
+
+Silently verify schema, string coordinates, Z-order, exact-text counts, geometry, bbox validity, and completeness.

code/diffusion/autoencoder_kl_qwenimage.py

@@ -0,0 +1,1057 @@
+#!/usr/bin/python
+# Copyright 2025 The Qwen-Image Team, Wan Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# We gratefully acknowledge the Wan Team for their outstanding contributions.
+# QwenImageVAE is further fine-tuned from the Wan Video VAE to achieve improved performance.
+# For more information about the Wan VAE, please refer to:
+# - GitHub: https://github.com/Wan-Video/Wan2.1
+# - Paper: https://huggingface.co/papers/2503.20314
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin
+from diffusers.utils import logging
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from diffusers.models.activations import get_activation
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.autoencoders.vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+CACHE_T = 2
+
+
+class QwenImageCausalConv3d(nn.Conv3d):
+    r"""
+    A custom 3D causal convolution layer with feature caching support.
+
+    This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature
+    caching for efficient inference.
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, int, int],
+        stride: int | tuple[int, int, int] = 1,
+        padding: int | tuple[int, int, int] = 0,
+    ) -> None:
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+        )
+
+        # Set up causal padding
+        self._padding = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0)
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+        return super().forward(x)
+
+
+class QwenImageRMS_norm(nn.Module):
+    r"""
+    A custom RMS normalization layer.
+
+    Args:
+        dim (int): The number of dimensions to normalize over.
+        channel_first (bool, optional): Whether the input tensor has channels as the first dimension.
+            Default is True.
+        images (bool, optional): Whether the input represents image data. Default is True.
+        bias (bool, optional): Whether to include a learnable bias term. Default is False.
+    """
+
+    def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any(
+            t in str(x.dtype) for t in ("float4_", "float8_")
+        )
+        normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=(1 if self.channel_first else -1)).to(
+            x.dtype
+        )
+
+        return normalized * self.scale * self.gamma + self.bias
+
+
+class QwenImageUpsample(nn.Upsample):
+    r"""
+    Perform upsampling while ensuring the output tensor has the same data type as the input.
+
+    Args:
+        x (torch.Tensor): Input tensor to be upsampled.
+
+    Returns:
+        torch.Tensor: Upsampled tensor with the same data type as the input.
+    """
+
+    def forward(self, x):
+        return super().forward(x.float()).type_as(x)
+
+
+class QwenImageResample(nn.Module):
+    r"""
+    A custom resampling module for 2D and 3D data.
+
+    Args:
+        dim (int): The number of input/output channels.
+        mode (str): The resampling mode. Must be one of:
+            - 'none': No resampling (identity operation).
+            - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution.
+            - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution.
+            - 'downsample2d': 2D downsampling with zero-padding and convolution.
+            - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution.
+    """
+
+    def __init__(self, dim: int, mode: str) -> None:
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                QwenImageUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                QwenImageUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+            self.time_conv = QwenImageCausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = QwenImageCausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == "upsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = "Rep"
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat(
+                            [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
+                        )
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
+                        cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+                    if feat_cache[idx] == "Rep":
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
+        x = self.resample(x)
+        x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
+
+        if self.mode == "downsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+
+class QwenImageResidualBlock(nn.Module):
+    r"""
+    A custom residual block module.
+
+    Args:
+        in_dim (int): Number of input channels.
+        out_dim (int): Number of output channels.
+        dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0.
+        non_linearity (str, optional): Type of non-linearity to use. Default is "silu".
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        dropout: float = 0.0,
+        non_linearity: str = "silu",
+    ) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.nonlinearity = get_activation(non_linearity)
+
+        # layers
+        self.norm1 = QwenImageRMS_norm(in_dim, images=False)
+        self.conv1 = QwenImageCausalConv3d(in_dim, out_dim, 3, padding=1)
+        self.norm2 = QwenImageRMS_norm(out_dim, images=False)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = QwenImageCausalConv3d(out_dim, out_dim, 3, padding=1)
+        self.conv_shortcut = QwenImageCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        # Apply shortcut connection
+        h = self.conv_shortcut(x)
+
+        # First normalization and activation
+        x = self.norm1(x)
+        x = self.nonlinearity(x)
+
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        # Second normalization and activation
+        x = self.norm2(x)
+        x = self.nonlinearity(x)
+
+        # Dropout
+        x = self.dropout(x)
+
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+
+            x = self.conv2(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv2(x)
+
+        # Add residual connection
+        return x + h
+
+
+class QwenImageAttentionBlock(nn.Module):
+    r"""
+    Causal self-attention with a single head.
+
+    Args:
+        dim (int): The number of channels in the input tensor.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = QwenImageRMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+    def forward(self, x):
+        identity = x
+        batch_size, channels, time, height, width = x.size()
+
+        x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width)
+        x = self.norm(x)
+
+        # compute query, key, value
+        qkv = self.to_qkv(x)
+        qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1)
+        qkv = qkv.permute(0, 1, 3, 2).contiguous()
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        # apply attention
+        x = F.scaled_dot_product_attention(q, k, v)
+
+        x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width)
+
+        # output projection
+        x = self.proj(x)
+
+        # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w]
+        x = x.view(batch_size, time, channels, height, width)
+        x = x.permute(0, 2, 1, 3, 4)
+
+        return x + identity
+
+
+class QwenImageMidBlock(nn.Module):
+    """
+    Middle block for QwenImageVAE encoder and decoder.
+
+    Args:
+        dim (int): Number of input/output channels.
+        dropout (float): Dropout rate.
+        non_linearity (str): Type of non-linearity to use.
+    """
+
+    def __init__(self, dim: int, dropout: float = 0.0, non_linearity: str = "silu", num_layers: int = 1):
+        super().__init__()
+        self.dim = dim
+
+        # Create the components
+        resnets = [QwenImageResidualBlock(dim, dim, dropout, non_linearity)]
+        attentions = []
+        for _ in range(num_layers):
+            attentions.append(QwenImageAttentionBlock(dim))
+            resnets.append(QwenImageResidualBlock(dim, dim, dropout, non_linearity))
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        # First residual block
+        x = self.resnets[0](x, feat_cache, feat_idx)
+
+        # Process through attention and residual blocks
+        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            if attn is not None:
+                x = attn(x)
+
+            x = resnet(x, feat_cache, feat_idx)
+
+        return x
+
+
+class QwenImageEncoder3d(nn.Module):
+    r"""
+    A 3D encoder module.
+
+    Args:
+        dim (int): The base number of channels in the first layer.
+        z_dim (int): The dimensionality of the latent space.
+        dim_mult (list of int): Multipliers for the number of channels in each block.
+        num_res_blocks (int): Number of residual blocks in each block.
+        attn_scales (list of float): Scales at which to apply attention mechanisms.
+        temperal_downsample (list of bool): Whether to downsample temporally in each block.
+        dropout (float): Dropout rate for the dropout layers.
+        non_linearity (str): Type of non-linearity to use.
+    """
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+        input_channels=3,
+        non_linearity: str = "silu",
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.nonlinearity = get_activation(non_linearity)
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv_in = QwenImageCausalConv3d(input_channels, dims[0], 3, padding=1)
+
+        # downsample blocks
+        self.down_blocks = nn.ModuleList([])
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            for _ in range(num_res_blocks):
+                self.down_blocks.append(QwenImageResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    self.down_blocks.append(QwenImageAttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
+                self.down_blocks.append(QwenImageResample(out_dim, mode=mode))
+                scale /= 2.0
+
+        # middle blocks
+        self.mid_block = QwenImageMidBlock(out_dim, dropout, non_linearity, num_layers=1)
+
+        # output blocks
+        self.norm_out = QwenImageRMS_norm(out_dim, images=False)
+        self.conv_out = QwenImageCausalConv3d(out_dim, z_dim, 3, padding=1)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+            x = self.conv_in(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv_in(x)
+
+        ## downsamples
+        for layer in self.down_blocks:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## middle
+        x = self.mid_block(x, feat_cache, feat_idx)
+
+        ## head
+        x = self.norm_out(x)
+        x = self.nonlinearity(x)
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+            x = self.conv_out(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv_out(x)
+        return x
+
+
+class QwenImageUpBlock(nn.Module):
+    """
+    A block that handles upsampling for the QwenImageVAE decoder.
+
+    Args:
+        in_dim (int): Input dimension
+        out_dim (int): Output dimension
+        num_res_blocks (int): Number of residual blocks
+        dropout (float): Dropout rate
+        upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d')
+        non_linearity (str): Type of non-linearity to use
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        num_res_blocks: int,
+        dropout: float = 0.0,
+        upsample_mode: str | None = None,
+        non_linearity: str = "silu",
+    ):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # Create layers list
+        resnets = []
+        # Add residual blocks and attention if needed
+        current_dim = in_dim
+        for _ in range(num_res_blocks + 1):
+            resnets.append(QwenImageResidualBlock(current_dim, out_dim, dropout, non_linearity))
+            current_dim = out_dim
+
+        self.resnets = nn.ModuleList(resnets)
+
+        # Add upsampling layer if needed
+        self.upsamplers = None
+        if upsample_mode is not None:
+            self.upsamplers = nn.ModuleList([QwenImageResample(out_dim, mode=upsample_mode)])
+
+        self.gradient_checkpointing = False
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        """
+        Forward pass through the upsampling block.
+
+        Args:
+            x (torch.Tensor): Input tensor
+            feat_cache (list, optional): Feature cache for causal convolutions
+            feat_idx (list, optional): Feature index for cache management
+
+        Returns:
+            torch.Tensor: Output tensor
+        """
+        for resnet in self.resnets:
+            if feat_cache is not None:
+                x = resnet(x, feat_cache, feat_idx)
+            else:
+                x = resnet(x)
+
+        if self.upsamplers is not None:
+            if feat_cache is not None:
+                x = self.upsamplers[0](x, feat_cache, feat_idx)
+            else:
+                x = self.upsamplers[0](x)
+        return x
+
+
+class QwenImageDecoder3d(nn.Module):
+    r"""
+    A 3D decoder module.
+
+    Args:
+        dim (int): The base number of channels in the first layer.
+        z_dim (int): The dimensionality of the latent space.
+        dim_mult (list of int): Multipliers for the number of channels in each block.
+        num_res_blocks (int): Number of residual blocks in each block.
+        attn_scales (list of float): Scales at which to apply attention mechanisms.
+        temperal_upsample (list of bool): Whether to upsample temporally in each block.
+        dropout (float): Dropout rate for the dropout layers.
+        non_linearity (str): Type of non-linearity to use.
+    """
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[False, True, True],
+        dropout=0.0,
+        input_channels=3,
+        non_linearity: str = "silu",
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        self.nonlinearity = get_activation(non_linearity)
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2 ** (len(dim_mult) - 2)
+
+        # init block
+        self.conv_in = QwenImageCausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.mid_block = QwenImageMidBlock(dims[0], dropout, non_linearity, num_layers=1)
+
+        # upsample blocks
+        self.up_blocks = nn.ModuleList([])
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            if i > 0:
+                in_dim = in_dim // 2
+
+            # Determine if we need upsampling
+            upsample_mode = None
+            if i != len(dim_mult) - 1:
+                upsample_mode = "upsample3d" if temperal_upsample[i] else "upsample2d"
+
+            # Create and add the upsampling block
+            up_block = QwenImageUpBlock(
+                in_dim=in_dim,
+                out_dim=out_dim,
+                num_res_blocks=num_res_blocks,
+                dropout=dropout,
+                upsample_mode=upsample_mode,
+                non_linearity=non_linearity,
+            )
+            self.up_blocks.append(up_block)
+
+            # Update scale for next iteration
+            if upsample_mode is not None:
+                scale *= 2.0
+
+        # output blocks
+        self.norm_out = QwenImageRMS_norm(out_dim, images=False)
+        self.conv_out = QwenImageCausalConv3d(out_dim, input_channels, 3, padding=1)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        ## conv1
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+            x = self.conv_in(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv_in(x)
+
+        ## middle
+        x = self.mid_block(x, feat_cache, feat_idx)
+
+        ## upsamples
+        for up_block in self.up_blocks:
+            x = up_block(x, feat_cache, feat_idx)
+
+        ## head
+        x = self.norm_out(x)
+        x = self.nonlinearity(x)
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+            x = self.conv_out(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv_out(x)
+        return x
+
+
+class AutoencoderKLQwenImage(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin):
+    r"""
+    A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+    """
+
+    _supports_gradient_checkpointing = False
+
+    # fmt: off
+    @register_to_config
+    def __init__(
+        self,
+        base_dim: int = 96,
+        z_dim: int = 16,
+        dim_mult: list[int] = [1, 2, 4, 4],
+        num_res_blocks: int = 2,
+        attn_scales: list[float] = [],
+        temperal_downsample: list[bool] = [False, True, True],
+        dropout: float = 0.0,
+        input_channels: int = 3,
+        latents_mean: list[float] = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921],
+        latents_std: list[float] = [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160],
+    ) -> None:
+    # fmt: on
+        super().__init__()
+
+        self.z_dim = z_dim
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        self.encoder = QwenImageEncoder3d(
+            base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout, input_channels
+        )
+        self.quant_conv = QwenImageCausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.post_quant_conv = QwenImageCausalConv3d(z_dim, z_dim, 1)
+
+        self.decoder = QwenImageDecoder3d(
+            base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout, input_channels
+        )
+
+        self.spatial_compression_ratio = 2 ** len(self.temperal_downsample)
+
+        # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension
+        # to perform decoding of a single video latent at a time.
+        self.use_slicing = False
+
+        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
+        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
+        # intermediate tiles together, the memory requirement can be lowered.
+        self.use_tiling = False
+
+        # The minimal tile height and width for spatial tiling to be used
+        self.tile_sample_min_height = 256
+        self.tile_sample_min_width = 256
+
+        # The minimal distance between two spatial tiles
+        self.tile_sample_stride_height = 192
+        self.tile_sample_stride_width = 192
+
+        # Precompute and cache conv counts for encoder and decoder for clear_cache speedup
+        self._cached_conv_counts = {
+            "decoder": sum(isinstance(m, QwenImageCausalConv3d) for m in self.decoder.modules())
+            if self.decoder is not None
+            else 0,
+            "encoder": sum(isinstance(m, QwenImageCausalConv3d) for m in self.encoder.modules())
+            if self.encoder is not None
+            else 0,
+        }
+
+    def enable_tiling(
+        self,
+        tile_sample_min_height: int | None = None,
+        tile_sample_min_width: int | None = None,
+        tile_sample_stride_height: float | None = None,
+        tile_sample_stride_width: float | None = None,
+    ) -> None:
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+
+        Args:
+            tile_sample_min_height (`int`, *optional*):
+                The minimum height required for a sample to be separated into tiles across the height dimension.
+            tile_sample_min_width (`int`, *optional*):
+                The minimum width required for a sample to be separated into tiles across the width dimension.
+            tile_sample_stride_height (`int`, *optional*):
+                The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are
+                no tiling artifacts produced across the height dimension.
+            tile_sample_stride_width (`int`, *optional*):
+                The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling
+                artifacts produced across the width dimension.
+        """
+        self.use_tiling = True
+        self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
+        self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
+        self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
+        self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
+
+    def clear_cache(self):
+        def _count_conv3d(model):
+            count = 0
+            for m in model.modules():
+                if isinstance(m, QwenImageCausalConv3d):
+                    count += 1
+            return count
+
+        self._conv_num = _count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        # cache encode
+        self._enc_conv_num = _count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+    def _encode(self, x: torch.Tensor):
+        _, _, num_frame, height, width = x.shape
+
+        if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height):
+            return self.tiled_encode(x)
+
+        self.clear_cache()
+        iter_ = 1 + (num_frame - 1) // 4
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+
+        enc = self.quant_conv(out)
+        self.clear_cache()
+        return enc
+
+    @apply_forward_hook
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
+        r"""
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded videos. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, z: torch.Tensor, return_dict: bool = True):
+        _, _, num_frame, height, width = z.shape
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+
+        if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
+            return self.tiled_decode(z, return_dict=return_dict)
+
+        self.clear_cache()
+        x = self.post_quant_conv(z)
+        for i in range(num_frame):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+            else:
+                out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2)
+
+        out = torch.clamp(out, min=-1.0, max=1.0)
+        self.clear_cache()
+        if not return_dict:
+            return (out,)
+
+        return DecoderOutput(sample=out)
+
+    @apply_forward_hook
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+        r"""
+        Decode a batch of images.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        if not return_dict:
+            return (decoded,)
+        return DecoderOutput(sample=decoded)
+
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
+                y / blend_extent
+            )
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
+                x / blend_extent
+            )
+        return b
+
+    def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        Args:
+            x (`torch.Tensor`): Input batch of videos.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+        _, _, num_frames, height, width = x.shape
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                self.clear_cache()
+                time = []
+                frame_range = 1 + (num_frames - 1) // 4
+                for k in range(frame_range):
+                    self._enc_conv_idx = [0]
+                    if k == 0:
+                        tile = x[:, :, :1, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + 4 * (k - 1) : 1 + 4 * k,
+                            i : i + self.tile_sample_min_height,
+                            j : j + self.tile_sample_min_width,
+                        ]
+                    tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+                    tile = self.quant_conv(tile)
+                    time.append(tile)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
+        self.clear_cache()
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
+        return enc
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, tile_latent_stride_height):
+            row = []
+            for j in range(0, width, tile_latent_stride_width):
+                self.clear_cache()
+                time = []
+                for k in range(num_frames):
+                    self._conv_idx = [0]
+                    tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                    tile = self.post_quant_conv(tile)
+                    decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                    time.append(decoded)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
+        self.clear_cache()
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
+
+        if not return_dict:
+            return (dec,)
+        return DecoderOutput(sample=dec)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: torch.Generator | None = None,
+    ) -> DecoderOutput | torch.Tensor:
+        """
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z, return_dict=return_dict)
+        return dec

code/diffusion/generator.py

@@ -0,0 +1,298 @@
+import torch
+from diffusers import AutoencoderKL
+import json
+import os
+from diffusers import FlowMatchEulerDiscreteScheduler
+from .transformer import DiffusionTransformer
+from .pipeline import ImageGenerationPipeline
+import torch.nn as nn
+import torch.nn.functional as F
+from .autoencoder_kl_qwenimage import AutoencoderKLQwenImage
+from inference_profile import InferenceProfile
+
+import logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class ToClipMLP(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        #self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(input_dim, 2048)
+        self.layer_norm1 = nn.LayerNorm(2048)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(2048, output_dim)
+        self.layer_norm2 = nn.LayerNorm(output_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states = self.relu(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = self.layer_norm2(hidden_states)
+        return hidden_states
+
+class ConditionedTransformer(nn.Module):
+    def __init__(self, transformer, vision_dim=1152, use_identity_mlp=False, text_encoder_norm=False):
+        super().__init__()
+        self.transformer = transformer
+        self.mlp = ToClipMLP(vision_dim, 2560) if not use_identity_mlp else nn.Identity()
+        self.mlp.to(dtype=self.dtype)
+        # self.mlp_pool = ToClipMLP(vision_dim, 768)
+        self.config = self.transformer.config
+        self.in_channels = self.transformer.in_channels
+        self.text_encoder_norm = text_encoder_norm
+
+        # must be used together
+        #if text_encoder_norm or use_identity_mlp:
+        #    assert use_identity_mlp and text_encoder_norm
+
+    
+    @property
+    def dtype(self):
+        return next(self.transformer.parameters()).dtype
+
+    def forward(self, hidden_states,
+                    timestep,
+                    encoder_hidden_states,
+                    return_dict,
+                    encoder_attention_mask=None,
+                    extra_vit_input=None,
+                    ref_hidden_states=None,
+                    encoder_hidden_states_2=None,
+                     **kargs):
+
+        if encoder_hidden_states is not None:
+            if isinstance(encoder_hidden_states, list):
+                encoder_hidden_states = torch.stack(encoder_hidden_states, dim=0)
+
+            if self.text_encoder_norm:
+                encoder_hidden_states = F.normalize(encoder_hidden_states, dim=-1) * 1000.0  # 1000 matches the original text encoder norm
+            
+            encoder_hidden_states = self.mlp(encoder_hidden_states)
+            
+            if extra_vit_input is not None:
+                encoder_hidden_states = torch.cat((encoder_hidden_states, extra_vit_input), dim=1)
+
+            encoder_hidden_states = list(encoder_hidden_states.unbind(dim=0))
+            
+        hidden_states = self.transformer(
+                    x=hidden_states,
+                    cap_feats=encoder_hidden_states,
+                    t=timestep,
+                    return_dict=False,
+                    ref_x=ref_hidden_states,
+                    cap_feats_2=encoder_hidden_states_2,
+                     **kargs
+                )
+        return hidden_states
+
+    def enable_gradient_checkpointing(self):
+        self.transformer.enable_gradient_checkpointing()
+
+def latent_mean_variance(latents: torch.Tensor, per_channel: bool = True, unbiased_var: bool = False):
+    """Compute the mean and variance of latents (shape follows the diffusion
+    training ``latents``, usually ``[B, C, H, W]``).
+
+    Args:
+        latents: ``[B, C, H, W]``; a ``[B, C, T, H, W]`` tensor with
+            ``T==1`` is ``squeeze(2)`` first.
+        per_channel: when True aggregate over ``(B, H, W)``, one scalar per
+            channel, shape ``[C]``; when False use one global mean/variance.
+        unbiased_var: use the unbiased estimator (Bessel).
+
+    Returns:
+        ``(mean, variance)``; the std is ``variance.sqrt()`` (use
+        ``torch.sqrt(variance.clamp_min(0))`` for numerical stability).
+    """
+    z = latents
+    if z.dim() == 5 and z.shape[2] == 1:
+        z = z.squeeze(2)
+    if z.dim() != 4:
+        raise ValueError(f"Expected 4D latents [B,C,H,W] (or 5D with T=1), got {tuple(z.shape)}")
+    if per_channel:
+        dims = (0, 2, 3)
+        mean_v = z.mean(dim=dims)
+        var_v = z.var(dim=dims, unbiased=unbiased_var)
+    else:
+        mean_v = z.mean()
+        var_v = z.var(unbiased=unbiased_var)
+    return mean_v, var_v
+
+class ImageGenerator(torch.nn.Module):
+    def __init__(self, 
+            model_path, 
+            vision_dim=2560, 
+            scheduler_path=None,
+            mlp_state_dict=None,
+            torch_dtype=torch.float32,
+            device='cpu',
+            use_identity_mlp=False,
+            text_encoder_norm=False,
+            inference_profile=None,
+        ):
+        super(ImageGenerator, self).__init__()
+
+        if not isinstance(inference_profile, InferenceProfile):
+            raise ValueError(
+                "inference_profile must be derived from the checkpoint "
+                "capability contract (load_checkpoint_capabilities)"
+            )
+        self.inference_profile = inference_profile
+
+        if device is not None:
+            device = torch.device(device)
+        else:
+            device = torch.device(torch.cuda.current_device())
+
+        self.scheduler_path = scheduler_path
+
+        vae_config_path = os.path.join(model_path, "vae", "config.json")
+        assert os.path.exists(vae_config_path)
+
+        with open(vae_config_path, "r") as f:
+            vae_config =  json.load(f)
+            if "_class_name" in vae_config and vae_config["_class_name"] == "AutoencoderKLQwenImage":
+                self.vae = AutoencoderKLQwenImage.from_pretrained(
+                    model_path,
+                    subfolder="vae",
+                    torch_dtype=torch_dtype,
+                )
+                self.vae_sample_mode = "argmax"
+            else:
+                self.vae = AutoencoderKL.from_pretrained(
+                    model_path,
+                    subfolder="vae",
+                    torch_dtype=torch_dtype,
+                )
+                self.vae_sample_mode = "sample"
+        
+        self.vae.input_channels = 4 if ('input_channels' in self.vae.config and self.vae.config.input_channels == 4) or ('in_channels' in self.vae.config and self.vae.config.in_channels == 4) else 3
+        if self.vae.input_channels != self.inference_profile.vae_input_channels:
+            raise ValueError(
+                "VAE input channels do not match the checkpoint capability "
+                f"contract: checkpoint={self.vae.input_channels}, "
+                f"capability={self.inference_profile.vae_input_channels}"
+            )
+        if self.vae_sample_mode != self.inference_profile.vae_sample_mode:
+            raise ValueError(
+                "VAE sample mode does not match the checkpoint capability "
+                f"contract: checkpoint={self.vae_sample_mode}, "
+                f"capability={self.inference_profile.vae_sample_mode}"
+            )
+        
+        # self.vae.to(self.torch_type).to(self.device)
+        self.vae.requires_grad_(False)
+
+        self.train_model = DiffusionTransformer.from_pretrained(
+            model_path, subfolder="transformer",
+            torch_dtype=torch_dtype,
+            alignment_padding_mode=self.inference_profile.alignment_padding_mode,
+            multi_frame_output=self.inference_profile.multi_frame_output,
+        )
+        if (
+            self.train_model.alignment_padding_mode
+            != self.inference_profile.alignment_padding_mode
+            or self.train_model.multi_frame_output
+            != self.inference_profile.multi_frame_output
+        ):
+            raise ValueError(
+                "instantiated Transformer capability does not match the "
+                "checkpoint capability contract: "
+                f"transformer=({self.train_model.alignment_padding_mode!r}, "
+                f"{self.train_model.multi_frame_output!r}), "
+                f"capability=({self.inference_profile.alignment_padding_mode!r}, "
+                f"{self.inference_profile.multi_frame_output!r})"
+            )
+
+        self.train_model = ConditionedTransformer(self.train_model, vision_dim=vision_dim, use_identity_mlp=use_identity_mlp, text_encoder_norm=text_encoder_norm)
+
+        assert mlp_state_dict is not None
+        self.train_model.mlp.load_state_dict(mlp_state_dict, strict=True)
+
+        self.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(self.scheduler_path, subfolder="scheduler")
+        self.noise_scheduler.config['use_dynamic_shifting'] = True
+
+        self.pipelines = ImageGenerationPipeline(
+            vae=self.vae,
+            transformer=self.train_model,  
+            text_encoder=None, 
+            tokenizer=None,                          
+            scheduler=self.noise_scheduler,
+        ).to(device)
+
+    @property
+    def device(self):
+        return next(self.train_model.parameters()).device
+
+    def set_trainable_params(self, trainable_params):
+        
+        self.vae.requires_grad_(False)
+
+        if trainable_params == 'all':
+            self.train_model.requires_grad_(True)
+        else:
+            self.train_model.requires_grad_(False)
+            for name, module in self.train_model.named_modules():
+                for trainable_param in trainable_params:
+                    if trainable_param in name:
+                        for params in module.parameters():
+                            params.requires_grad = True
+
+        num_parameters_trainable = 0
+        num_parameters = 0
+        name_parameters_trainable = []
+        for n, p in self.train_model.named_parameters():
+            num_parameters += p.data.nelement()
+            if not p.requires_grad:
+                continue  # frozen weights
+            name_parameters_trainable.append(n)
+            num_parameters_trainable += p.data.nelement()
+        logger.info(f"number of all Diffusion parameters: {num_parameters}, trainable: {num_parameters_trainable}")
+    
+
+    def sample(self, encoder_hidden_states, steps=None, cfg=None, cfg_mode=1, seed=42, height=512, width=512, use_dynamic_shifting=False, extra_vit_input=None, ref_x=None, directvlm_hidden_states=None, num_frames_per_prompt=1):
+        sampling = self.inference_profile.resolve_sampling_parameters(
+            steps=steps,
+            cfg=cfg,
+        )
+        steps = sampling.steps
+        cfg = sampling.cfg
+        negative_prompt_embeds = None
+        if encoder_hidden_states is not None:
+            encoder_hidden_states = encoder_hidden_states.to(
+                device=self.device, dtype=self.train_model.dtype
+            )
+            encoder_hidden_states = list(encoder_hidden_states.unbind(dim=0))
+            negative_prompt_embeds= [en * 0 for en in encoder_hidden_states]
+        
+        encoder_hidden_states_2 = directvlm_hidden_states
+        negative_prompt_embeds_2 = None
+        if encoder_hidden_states_2 is not None:
+            encoder_hidden_states_2 = encoder_hidden_states_2.to(
+                device=self.device, dtype=self.train_model.dtype
+            )
+            encoder_hidden_states_2 = list(encoder_hidden_states_2.unbind(dim=0))
+            negative_prompt_embeds_2= [en * 0 for en in encoder_hidden_states_2]
+
+        image = self.pipelines(
+            prompt_embeds=encoder_hidden_states,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_embeds_2=encoder_hidden_states_2,
+            negative_prompt_embeds_2=negative_prompt_embeds_2,
+            guidance_scale=cfg,
+            #guidance_scale_mode=cfg_mode,
+            generator=torch.manual_seed(seed),
+            num_inference_steps=steps,
+            height=height,
+            width=width,
+            max_sequence_length=512,
+            device=self.device,
+            #extra_vit_input=extra_vit_input,
+            ref_hidden_states=ref_x,
+            #use_dynamic_shifting=use_dynamic_shifting,
+            sample_mode=self.vae_sample_mode,
+            num_frames_per_prompt=num_frames_per_prompt,
+        ).images
+
+        return image  

code/diffusion/padding.py

@@ -0,0 +1,45 @@
+"""Alignment-padding helpers for the standard diffusion transformer."""
+
+from typing import Optional, Sequence
+
+import torch
+
+
+def mask_out_alignment_padding(
+    attention_mask: torch.Tensor,
+    pad_masks: Sequence[Optional[torch.Tensor]],
+    offsets: Sequence[int],
+) -> torch.Tensor:
+    """Exclude per-item alignment padding from a 2D boolean attention mask."""
+
+    if attention_mask.ndim != 2 or attention_mask.dtype != torch.bool:
+        raise ValueError(
+            "attention_mask must be a 2D boolean tensor, got "
+            f"shape={attention_mask.shape}, dtype={attention_mask.dtype}"
+        )
+    batch_size = attention_mask.shape[0]
+    if len(pad_masks) != batch_size or len(offsets) != batch_size:
+        raise ValueError(
+            "pad mask metadata must match the attention-mask batch size: "
+            f"batch={batch_size}, pad_masks={len(pad_masks)}, offsets={len(offsets)}"
+        )
+
+    for item_index, (pad_mask, offset) in enumerate(zip(pad_masks, offsets)):
+        if pad_mask is None or pad_mask.numel() == 0:
+            continue
+        if pad_mask.ndim != 1 or pad_mask.dtype != torch.bool:
+            raise ValueError(
+                "each alignment-pad mask must be a 1D boolean tensor: "
+                f"item={item_index}, shape={pad_mask.shape}, dtype={pad_mask.dtype}"
+            )
+        offset = int(offset)
+        end = offset + pad_mask.shape[0]
+        if offset < 0 or end > attention_mask.shape[1]:
+            raise ValueError(
+                "alignment-pad mask falls outside the attention sequence: "
+                f"item={item_index}, offset={offset}, length={pad_mask.shape[0]}, "
+                f"sequence={attention_mask.shape[1]}"
+            )
+        attention_mask[item_index, offset:end].masked_fill_(pad_mask, False)
+
+    return attention_mask

code/diffusion/pipeline.py

@@ -0,0 +1,716 @@
+# Copyright (c) Ant Group. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# All rights and credit for the original implementation remain with the original authors and contributors, and this project complies with the applicable open-source license terms of the referenced repository.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from transformers import AutoTokenizer, PreTrainedModel
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+
+from .transformer import DiffusionTransformer
+
+from dataclasses import dataclass
+from typing import List, Union
+
+import numpy as np
+import PIL.Image
+
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class ImageGenerationOutput(BaseOutput):
+    """
+    Output class for image generation pipelines.
+
+    Args:
+        images (`List[PIL.Image.Image]` or `np.ndarray`)
+            List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+            num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
+    """
+
+    images: Union[List[PIL.Image.Image], np.ndarray]
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        This pipeline is assembled by the Ming image checkpoint loader; the
+        supported entry point is the repository CLI:
+
+        ```bash
+        python infer.py \
+          --model /path/to/checkpoint-package \
+          --task text-to-image \
+          --prompt "A small red cube on a white table" \
+          --output-dir outputs/t2i
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class ImageGenerationPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: PreTrainedModel,
+        tokenizer: AutoTokenizer,
+        transformer: DiffusionTransformer,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            scheduler=scheduler,
+            transformer=transformer,
+        )
+
+        if hasattr(self, "vae") and self.vae is not None and 'temperal_downsample' in self.vae.config:
+            self.vae_scale_factor = 2 ** len(self.vae.config.temperal_downsample)
+        else:
+            self.vae_scale_factor = (
+                2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
+            )
+
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        max_sequence_length: int = 512,
+    ):
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prompt_embeds = self._encode_prompt(
+            prompt=prompt,
+            device=device,
+            prompt_embeds=prompt_embeds,
+            max_sequence_length=max_sequence_length,
+        )
+
+        if do_classifier_free_guidance:
+            if negative_prompt is None:
+                negative_prompt = ["" for _ in prompt]
+            else:
+                negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            assert len(prompt) == len(negative_prompt)
+            negative_prompt_embeds = self._encode_prompt(
+                prompt=negative_prompt,
+                device=device,
+                prompt_embeds=negative_prompt_embeds,
+                max_sequence_length=max_sequence_length,
+            )
+        else:
+            negative_prompt_embeds = []
+        return prompt_embeds, negative_prompt_embeds
+
+    def _encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+        max_sequence_length: int = 512,
+    ) -> List[torch.FloatTensor]:
+        device = device or self._execution_device
+
+        if prompt_embeds is not None:
+            return prompt_embeds
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        for i, prompt_item in enumerate(prompt):
+            messages = [
+                {"role": "user", "content": prompt_item},
+            ]
+            prompt_item = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=True,
+            )
+            prompt[i] = prompt_item
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids.to(device)
+        prompt_masks = text_inputs.attention_mask.to(device).bool()
+
+        prompt_embeds = self.text_encoder(
+            input_ids=text_input_ids,
+            attention_mask=prompt_masks,
+            output_hidden_states=True,
+        ).hidden_states[-2]
+
+        embeddings_list = []
+
+        for i in range(len(prompt_embeds)):
+            embeddings_list.append(prompt_embeds[i][prompt_masks[i]])
+
+        return embeddings_list
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            if latents.shape != shape:
+                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+            latents = latents.to(device)
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 5.0,
+        cfg_normalization: bool = False,
+        cfg_truncation: float = 1.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        num_frames_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+        negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+        device: Optional[Union[str, torch.device]] = None,
+        ref_hidden_states: Optional[torch.FloatTensor] = None,
+        prompt_embeds_2: Optional[List[torch.FloatTensor]] = None,
+        negative_prompt_embeds_2: Optional[List[torch.FloatTensor]] = None,
+        sample_mode: Optional[str] = "sample",
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`int`, *optional*, defaults to 1024):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to 1024):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 5.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            cfg_normalization (`bool`, *optional*, defaults to False):
+                Whether to apply configuration normalization.
+            cfg_truncation (`float`, *optional*, defaults to 1.0):
+                The truncation value for configuration.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will be generated by sampling using the supplied random `generator`.
+            prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.ImageGenerationOutput`] instead of a plain
+                tuple.
+            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).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int`, *optional*, defaults to 512):
+                Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`ImageGenerationOutput`] or `tuple`: [`ImageGenerationOutput`] if
+            `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
+            generated images.
+        """
+        height = height or 1024
+        width = width or 1024
+        if type(num_frames_per_prompt) is not int or num_frames_per_prompt < 1:
+            raise ValueError("num_frames_per_prompt must be an integer >= 1")
+        supports_multi_frame = bool(
+            getattr(self.transformer.config, "multi_frame_output", False)
+        )
+        if num_frames_per_prompt > 1 and not supports_multi_frame:
+            raise ValueError(
+                "the loaded checkpoint profile does not support multi-frame output"
+            )
+
+        vae_scale = self.vae_scale_factor * 2
+        if height % vae_scale != 0:
+            raise ValueError(
+                f"Height must be divisible by {vae_scale} (got {height}). "
+                f"Please adjust the height to a multiple of {vae_scale}."
+            )
+        if width % vae_scale != 0:
+            raise ValueError(
+                f"Width must be divisible by {vae_scale} (got {width}). "
+                f"Please adjust the width to a multiple of {vae_scale}."
+            )
+
+        device = device or self._execution_device
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+        self._cfg_normalization = cfg_normalization
+        self._cfg_truncation = cfg_truncation
+        # 2. Define call parameters
+
+        self.vae.config.shift_factor = self.vae.config.shift_factor.to(device=device, dtype=self.vae.dtype) if isinstance(self.vae.config.shift_factor, torch.Tensor) else self.vae.config.shift_factor
+        self.vae.config.scaling_factor = self.vae.config.scaling_factor.to(device=device, dtype=self.vae.dtype) if isinstance(self.vae.config.scaling_factor, torch.Tensor) else self.vae.config.scaling_factor
+
+        assert prompt is None
+        assert prompt_embeds is not None or prompt_embeds_2 is not None
+        if prompt_embeds is not None:
+            batch_size = len(prompt_embeds)
+            if prompt_embeds_2 is not None:
+                assert len(prompt_embeds_2) == len(prompt_embeds)
+
+            assert negative_prompt_embeds is not None
+            assert len(negative_prompt_embeds) == len(prompt_embeds)
+        else:
+            assert prompt_embeds_2 is not None
+            batch_size = len(prompt_embeds_2)
+            assert negative_prompt_embeds_2 is not None
+            assert len(negative_prompt_embeds_2) == len(prompt_embeds_2)
+            
+
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.in_channels
+
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt * num_frames_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            torch.float32,
+            device,
+            generator,
+            latents,
+        )
+
+        if ref_hidden_states is not None:
+            assert ref_hidden_states.ndim == 4  # and ref_hidden_states.shape[0] == 1
+            ref_hidden_states = ref_hidden_states.to(self.vae.dtype).to(device)
+            C = ref_hidden_states.shape[1]
+            input_channels = self.vae.input_channels
+            if C % input_channels != 0:
+               raise ValueError(f"ref_hidden_states.shape[1] must be a multiple of {input_channels}, got C={C}")
+            num_images = C // input_channels
+            ref_latents = []
+            for i in range(num_images):
+                x = ref_hidden_states[:, i * input_channels:(i + 1) * input_channels, :, :]  # RGB block of image i
+                if 'temperal_downsample' in self.vae.config:
+                    x = x.unsqueeze(2)
+                if sample_mode == "argmax":
+                    z = self.vae.encode(x).latent_dist.mode()
+                elif sample_mode == "sample":
+                    z = self.vae.encode(x).latent_dist.sample()
+                z = (z - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+                ref_latents.append(z)
+            
+            ref_hidden_states = torch.cat(ref_latents, dim=1)
+            
+            if ref_hidden_states.dim() == 5:
+                ref_hidden_states = ref_hidden_states.squeeze(2)
+
+
+
+        # Repeat prompt_embeds for num_images_per_prompt
+        if num_images_per_prompt > 1:
+            if prompt_embeds is not None:
+                prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)]
+                if self.do_classifier_free_guidance and negative_prompt_embeds:
+                    negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)]
+
+            if prompt_embeds_2 is not None:
+                prompt_embeds_2 = [pe for pe in prompt_embeds_2 for _ in range(num_images_per_prompt)]
+                if self.do_classifier_free_guidance and negative_prompt_embeds_2:
+                    negative_prompt_embeds_2 = [npe for npe in negative_prompt_embeds_2 for _ in range(num_images_per_prompt)]
+
+        actual_batch_size = batch_size * num_images_per_prompt
+        image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
+
+        # 5. Prepare timesteps
+        if image_seq_len >= 4096:
+            self.scheduler.config["max_image_seq_len"] = image_seq_len
+            self.scheduler.config["max_shift"] = 1.35
+        else:
+            self.scheduler.config["max_image_seq_len"] = 4096
+            self.scheduler.config["max_shift"] = 1.15
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        self.scheduler.sigma_min = 0.0
+        scheduler_kwargs = {"mu": mu}
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            **scheduler_kwargs,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        if num_frames_per_prompt > 1:
+            latents = torch.stack(
+                latents.chunk(num_frames_per_prompt, dim=0), dim=2
+            )
+
+        # 6. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0])
+                timestep = (1000 - timestep) / 1000
+                # Normalized time for time-aware config (0 at start, 1 at end)
+                t_norm = timestep[0].item()
+
+                # Handle cfg truncation
+                current_guidance_scale = self.guidance_scale
+                if (
+                    self.do_classifier_free_guidance
+                    and self._cfg_truncation is not None
+                    and float(self._cfg_truncation) <= 1
+                ):
+                    if t_norm > self._cfg_truncation:
+                        current_guidance_scale = 0.0
+
+                # Run CFG only if configured AND scale is non-zero
+                apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0
+
+                if apply_cfg:
+                    latents_typed = latents.to(self.transformer.dtype)
+                    repeat_dims = (2, 1, 1, 1, 1) if latents_typed.ndim == 5 else (2, 1, 1, 1)
+                    latent_model_input = latents_typed.repeat(*repeat_dims)
+                    prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds if prompt_embeds is not None else None
+                    prompt_embeds_model_input_2 = prompt_embeds_2 + negative_prompt_embeds_2 if prompt_embeds_2 is not None else None
+                    timestep_model_input = timestep.repeat(2)
+                    ref_hidden_states_input = ref_hidden_states.repeat(2, 1, 1, 1) if ref_hidden_states is not None else None
+                    if ref_hidden_states_input is not None:
+                        ref_hidden_states_input = ref_hidden_states_input.to(latent_model_input.dtype)
+                else:
+                    latent_model_input = latents.to(self.transformer.dtype)
+                    prompt_embeds_model_input = prompt_embeds
+                    prompt_embeds_model_input_2 = prompt_embeds_2
+                    timestep_model_input = timestep
+                    ref_hidden_states_input = ref_hidden_states*1.0 if ref_hidden_states is not None else None
+                    if ref_hidden_states_input is not None:
+                        ref_hidden_states_input = ref_hidden_states_input.to(latent_model_input.dtype)
+
+                if latent_model_input.ndim == 4:
+                    latent_model_input = latent_model_input.unsqueeze(2)
+                latent_model_input_list = list(latent_model_input.unbind(dim=0))
+                
+                if ref_hidden_states_input is not None:
+                    C = self.vae.config.latent_channels if 'latent_channels' in self.vae.config else self.vae.config.z_dim
+
+                    # single-image input: [B, C, H, W] -> list of B elems, each [C, 1, H, W]
+                    if ref_hidden_states_input.shape[1] == C:
+                        ref_hidden_states_input = ref_hidden_states_input.unsqueeze(2)              # [B, C, 1, H, W]
+                        ref_hidden_states_input = list(ref_hidden_states_input.unbind(dim=0))       # B * [C, 1, H, W]
+
+                    # multi-image input: [B, N*C, H, W] -> list of B elems, each [C, N, H, W]
+                    else:
+                        B, total_C, H, W = ref_hidden_states_input.shape
+                        if total_C % C != 0:
+                            raise ValueError(
+                                f"ref_hidden_states_input channel ({total_C}) must be divisible by latent_channels ({C})."
+                            )
+                        N = total_C // C
+
+                        # order-preserving: input assumes [img0(C), img1(C), ..., imgN-1(C)] stacked along channels
+                        ref_hidden_states_input = (
+                            ref_hidden_states_input.reshape(B, N, C, H, W)         # [B, N, C, H, W], N in stacking order
+                                .permute(0, 2, 1, 3, 4)         # [B, C, N, H, W]
+                                .contiguous()
+                        )
+                        ref_hidden_states_input = list(ref_hidden_states_input.unbind(dim=0))        # B * [C, N, H, W]
+
+                model_out_list = self.transformer(
+                    latent_model_input_list,
+                    timestep_model_input,
+                    prompt_embeds_model_input,
+                    ref_hidden_states=ref_hidden_states_input,
+                    return_dict=False,
+                    encoder_hidden_states_2=prompt_embeds_model_input_2,
+                )[0]
+                model_out_list = [
+                    output[:, :num_frames_per_prompt, :, :].float()
+                    for output in model_out_list
+                ]
+
+                if apply_cfg:
+                    # Perform CFG
+                    pos_out = model_out_list[:actual_batch_size]
+                    neg_out = model_out_list[actual_batch_size:]
+
+                    noise_pred = []
+                    for j in range(actual_batch_size):
+                        pos = pos_out[j].float()
+                        neg = neg_out[j].float()
+
+                        pred = pos + current_guidance_scale * (pos - neg)
+
+                        # Renormalization
+                        if self._cfg_normalization and float(self._cfg_normalization) > 0.0:
+                            ori_pos_norm = torch.linalg.vector_norm(pos)
+                            new_pos_norm = torch.linalg.vector_norm(pred)
+                            max_new_norm = ori_pos_norm * float(self._cfg_normalization)
+                            if new_pos_norm > max_new_norm:
+                                pred = pred * (max_new_norm / new_pos_norm)
+
+                        noise_pred.append(pred)
+
+                    noise_pred = torch.stack(noise_pred, dim=0)
+                else:
+                    noise_pred = torch.stack([t.float() for t in model_out_list], dim=0)
+
+                if num_frames_per_prompt == 1:
+                    noise_pred = noise_pred.squeeze(2)
+                noise_pred = -noise_pred
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0]
+                assert latents.dtype == torch.float32
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+        if num_frames_per_prompt > 1:
+            latents = torch.cat(
+                latents.chunk(num_frames_per_prompt, dim=2), dim=0
+            ).squeeze(2)
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = latents.to(self.vae.dtype)
+            if 'temperal_downsample' in self.vae.config:
+                latents = latents.unsqueeze(2)
+
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            if image.dim() == 5:
+                image = image.squeeze(2)
+                
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return ImageGenerationOutput(images=image)

code/diffusion/transformer.py

@@ -0,0 +1,768 @@
+# Copyright (c) Ant Group. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention_processor import Attention
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import RMSNorm
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+from inference_profile import LEARNED_PADDING, VALID_PADDING_MODES, ZERO_MASKED_PADDING
+from .padding import mask_out_alignment_padding
+
+
+ADALN_EMBED_DIM = 256
+SEQ_MULTI_OF = 32
+
+
+def _native_sdpa_is_active(processor) -> bool:
+    """True when attention would go to diffusers' default native SDPA backend (no per-model backend,
+    no context parallelism, and the active global backend is NATIVE)."""
+    if processor._attention_backend is not None or processor._parallel_config is not None:
+        return False
+    try:
+        from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
+
+        name, _ = _AttentionBackendRegistry.get_active_backend()
+        return name == AttentionBackendName.NATIVE
+    except Exception:
+        return False
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, out_size, mid_size=None, frequency_embedding_size=256):
+        super().__init__()
+        if mid_size is None:
+            mid_size = out_size
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, mid_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(mid_size, out_size, bias=True),
+        )
+
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        with torch.amp.autocast("cuda", enabled=False):
+            half = dim // 2
+            freqs = torch.exp(
+                -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
+            )
+            args = t[:, None].float() * freqs[None]
+            embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+            if dim % 2:
+                embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+            return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        weight_dtype = self.mlp[0].weight.dtype
+        if weight_dtype.is_floating_point:
+            t_freq = t_freq.to(weight_dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class SingleStreamAttentionProcessor:
+    """
+    Processor for single-stream (joint text-image) attention; adapts the
+    existing Attention class to the single-stream behavior.
+    """
+
+    _attention_backend = None
+    _parallel_config = None
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "SingleStreamAttentionProcessor requires PyTorch 2.0. To use it, please upgrade PyTorch to version 2.0 or higher."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        freqs_cis: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        query = query.unflatten(-1, (attn.heads, -1))
+        key = key.unflatten(-1, (attn.heads, -1))
+        value = value.unflatten(-1, (attn.heads, -1))
+
+        # Apply Norms
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply RoPE
+        def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+            with torch.amp.autocast("cuda", enabled=False):
+                x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2))
+                freqs_cis = freqs_cis.unsqueeze(2)
+                x_out = torch.view_as_real(x * freqs_cis).flatten(3)
+                return x_out.type_as(x_in)  # todo
+
+        if freqs_cis is not None:
+            query = apply_rotary_emb(query, freqs_cis)
+            key = apply_rotary_emb(key, freqs_cis)
+
+        # Cast to correct dtype
+        dtype = query.dtype
+        query, key = query.to(dtype), key.to(dtype)
+
+        # From [batch, seq_len] to [batch, 1, 1, seq_len] -> broadcast to [batch, heads, seq_len, seq_len]
+        if attention_mask is not None and attention_mask.ndim == 2:
+            attention_mask = attention_mask[:, None, None, :]
+
+        # Compute joint attention
+        if _native_sdpa_is_active(self):
+            # What diffusers' default "native" backend computes, but SDPA receives contiguous
+            # [B, H, L, D] tensors instead of permuted views. PyTorch's math SDPA (the only SDPA
+            # kernel that runs on ROCm gfx1151) is ~2x faster on contiguous inputs, with
+            # bit-identical output.
+            q, k, v = (x.transpose(1, 2).contiguous() for x in (query, key, value))
+            hidden_states = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            ).transpose(1, 2)
+        else:
+            hidden_states = dispatch_attention_fn(
+                query,
+                key,
+                value,
+                attn_mask=attention_mask,
+                dropout_p=0.0,
+                is_causal=False,
+                backend=self._attention_backend,
+                parallel_config=self._parallel_config,
+            )
+
+        # Reshape back
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(dtype)
+
+        output = attn.to_out[0](hidden_states)
+        if len(attn.to_out) > 1:  # dropout
+            output = attn.to_out[1](output)
+
+        return output
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def _forward_silu_gating(self, x1, x3):
+        return F.silu(x1) * x3
+
+    def forward(self, x):
+        return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x)))
+
+
+@maybe_allow_in_graph
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        layer_id: int,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        norm_eps: float,
+        qk_norm: bool,
+        modulation=True,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.head_dim = dim // n_heads
+
+        # Refactored to use diffusers Attention with custom processor
+        # Upstream parameter order: dim, n_heads, n_kv_heads, qk_norm
+        self.attention = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=dim // n_heads,
+            heads=n_heads,
+            qk_norm="rms_norm" if qk_norm else None,
+            eps=1e-5,
+            bias=False,
+            out_bias=False,
+            processor=SingleStreamAttentionProcessor(),
+        )
+
+        self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8))
+        self.layer_id = layer_id
+
+        self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+
+        self.attention_norm2 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
+
+        self.modulation = modulation
+        if modulation:
+            self.adaLN_modulation = nn.Sequential(nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        adaln_input: Optional[torch.Tensor] = None,
+    ):
+        if self.modulation:
+            assert adaln_input is not None
+            scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2)
+            gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
+            scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
+
+            # Attention block
+            attn_out = self.attention(
+                self.attention_norm1(x) * scale_msa, attention_mask=attn_mask, freqs_cis=freqs_cis
+            )
+            x = x + gate_msa * self.attention_norm2(attn_out)
+
+            # FFN block
+            x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))
+        else:
+            # Attention block
+            attn_out = self.attention(self.attention_norm1(x), attention_mask=attn_mask, freqs_cis=freqs_cis)
+            x = x + self.attention_norm2(attn_out)
+
+            # FFN block
+            x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x)))
+
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True),
+        )
+
+    def forward(self, x, c):
+        scale = 1.0 + self.adaLN_modulation(c)
+        x = self.norm_final(x) * scale.unsqueeze(1)
+        x = self.linear(x)
+        return x
+
+
+class RopeEmbedder:
+    def __init__(
+        self,
+        theta: float = 256.0,
+        axes_dims: List[int] = (16, 56, 56),
+        axes_lens: List[int] = (64, 128, 128),
+    ):
+        self.theta = theta
+        self.axes_dims = axes_dims
+        self.axes_lens = axes_lens
+        assert len(axes_dims) == len(axes_lens), "axes_dims and axes_lens must have the same length"
+        self.freqs_cis = None
+
+    @staticmethod
+    def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0):
+        with torch.device("cpu"):
+            freqs_cis = []
+            for i, (d, e) in enumerate(zip(dim, end)):
+                freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d))
+                timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
+                freqs = torch.outer(timestep, freqs).float()
+                freqs_cis_i = torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64)  # complex64
+                freqs_cis.append(freqs_cis_i)
+
+            return freqs_cis
+
+    def __call__(self, ids: torch.Tensor):
+        assert ids.ndim == 2
+        assert ids.shape[-1] == len(self.axes_dims)
+        device = ids.device
+
+        if self.freqs_cis is None:
+            self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta)
+            self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+        else:
+            # Ensure freqs_cis are on the same device as ids
+            if self.freqs_cis[0].device != device:
+                self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+
+        result = []
+        for i in range(len(self.axes_dims)):
+            index = ids[:, i]
+            result.append(self.freqs_cis[i][index])
+        return torch.cat(result, dim=-1)
+
+
+class DiffusionTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["TransformerBlock"]
+    _repeated_blocks = ["TransformerBlock"]
+    _skip_layerwise_casting_patterns = ["t_embedder", "cap_embedder"]  # precision sensitive layers
+
+    @register_to_config
+    def __init__(
+        self,
+        all_patch_size=(2,),
+        all_f_patch_size=(1,),
+        in_channels=16,
+        dim=3840,
+        n_layers=30,
+        n_refiner_layers=2,
+        n_heads=30,
+        n_kv_heads=30,
+        norm_eps=1e-5,
+        qk_norm=True,
+        cap_feat_dim=2560,
+        rope_theta=256.0,
+        t_scale=1000.0,
+        axes_dims=[32, 48, 48],
+        axes_lens=[20480, 512, 512],
+        alignment_padding_mode=None,
+        multi_frame_output=None,
+    ) -> None:
+        super().__init__()
+        if alignment_padding_mode not in VALID_PADDING_MODES:
+            raise ValueError(
+                "alignment_padding_mode must come from the checkpoint capability "
+                "contract (transformer/config.json, or the legacy "
+                f"inference_profile.json) and be one of "
+                f"{sorted(VALID_PADDING_MODES)}, got {alignment_padding_mode!r}"
+            )
+        if type(multi_frame_output) is not bool:
+            raise ValueError(
+                "multi_frame_output must come from the checkpoint capability "
+                "contract (transformer/config.json, or the legacy "
+                "inference_profile.json) as a boolean"
+            )
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.all_patch_size = all_patch_size
+        self.all_f_patch_size = all_f_patch_size
+        self.dim = dim
+        self.n_heads = n_heads
+        self.alignment_padding_mode = alignment_padding_mode
+        self.multi_frame_output = multi_frame_output
+
+        self.rope_theta = rope_theta
+        self.t_scale = t_scale
+        self.gradient_checkpointing = False
+
+        assert len(all_patch_size) == len(all_f_patch_size)
+
+        all_x_embedder = {}
+        all_final_layer = {}
+        for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)):
+            x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * in_channels, dim, bias=True)
+            all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder
+
+            final_layer = FinalLayer(dim, patch_size * patch_size * f_patch_size * self.out_channels)
+            all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer
+
+        self.all_x_embedder = nn.ModuleDict(all_x_embedder)
+        self.all_final_layer = nn.ModuleDict(all_final_layer)
+        self.noise_refiner = nn.ModuleList(
+            [
+                TransformerBlock(
+                    1000 + layer_id,
+                    dim,
+                    n_heads,
+                    n_kv_heads,
+                    norm_eps,
+                    qk_norm,
+                    modulation=True,
+                )
+                for layer_id in range(n_refiner_layers)
+            ]
+        )
+        self.context_refiner = nn.ModuleList(
+            [
+                TransformerBlock(
+                    layer_id,
+                    dim,
+                    n_heads,
+                    n_kv_heads,
+                    norm_eps,
+                    qk_norm,
+                    modulation=False,
+                )
+                for layer_id in range(n_refiner_layers)
+            ]
+        )
+        self.t_embedder = TimestepEmbedder(min(dim, ADALN_EMBED_DIM), mid_size=1024)
+        self.cap_embedder = nn.Sequential(RMSNorm(cap_feat_dim, eps=norm_eps), nn.Linear(cap_feat_dim, dim, bias=True))
+
+        if self.alignment_padding_mode == LEARNED_PADDING:
+            self.x_pad_token = nn.Parameter(torch.empty((1, dim)))
+            self.cap_pad_token = nn.Parameter(torch.empty((1, dim)))
+        else:
+            self.register_parameter("x_pad_token", None)
+            self.register_parameter("cap_pad_token", None)
+
+        self.layers = nn.ModuleList(
+            [
+                TransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm)
+                for layer_id in range(n_layers)
+            ]
+        )
+        head_dim = dim // n_heads
+        assert head_dim == sum(axes_dims)
+        self.axes_dims = axes_dims
+        self.axes_lens = axes_lens
+
+        self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens)
+
+    def unpatchify(self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size) -> List[torch.Tensor]:
+        pH = pW = patch_size
+        pF = f_patch_size
+        bsz = len(x)
+        assert len(size) == bsz
+        for i in range(bsz):
+            F, H, W = size[i]
+            ori_len = (F // pF) * (H // pH) * (W // pW)
+            # "f h w pf ph pw c -> c (f pf) (h ph) (w pw)"
+            x[i] = (
+                x[i][:ori_len]
+                .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels)
+                .permute(6, 0, 3, 1, 4, 2, 5)
+                .reshape(self.out_channels, F, H, W)
+            )
+            if not self.multi_frame_output:
+                x[i] = x[i][:, :1, :, :]
+        return x
+
+    @staticmethod
+    def create_coordinate_grid(size, start=None, device=None):
+        if start is None:
+            start = (0 for _ in size)
+
+        axes = [torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size)]
+        grids = torch.meshgrid(axes, indexing="ij")
+        return torch.stack(grids, dim=-1)
+
+    def patchify_and_embed(
+        self,
+        all_image: List[torch.Tensor],
+        all_cap_feats: List[torch.Tensor],
+        patch_size: int,
+        f_patch_size: int,
+        all_image_ref: List[torch.Tensor] = None,
+        all_cap_feats_2: List[torch.Tensor] = None,
+    ):
+        pH = pW = patch_size
+        pF = f_patch_size
+        device = all_image[0].device
+
+        all_image_out = []
+        all_image_size = []
+        all_image_pos_ids = []
+        all_image_pad_mask = []
+        all_cap_pos_ids = []
+        all_cap_pad_mask = []
+        all_cap_feats_out = []
+        all_cap_feats_2_out = []
+
+        if all_image_ref is None:
+            all_image_ref = [None]*len(all_image)
+
+        assert all_cap_feats is not None or all_cap_feats_2 is not None
+
+        if all_cap_feats is None:
+            all_cap_feats = [None for _ in range(len(all_image))]
+        else:
+            assert not any([i is None for i in all_cap_feats])
+        
+        if all_cap_feats_2 is None:
+            all_cap_feats_2 = [None for _ in range(len(all_image))]
+        else:
+            assert not any([i is None for i in all_cap_feats_2])
+
+        for i, (image, cap_feat, cap_feat_2, image_ref) in enumerate(zip(all_image, all_cap_feats, all_cap_feats_2, all_image_ref)):
+            ### Process Caption
+            cap_ori_len = 0
+            if cap_feat is not None:
+                cap_ori_len += len(cap_feat)
+            
+            if cap_feat_2 is not None:
+                cap_ori_len += len(cap_feat_2)
+            
+            cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF
+            # padded position ids
+            cap_padded_pos_ids = self.create_coordinate_grid(
+                size=(cap_ori_len + cap_padding_len, 1, 1),
+                start=(1, 0, 0),
+                device=device,
+            ).flatten(0, 2)
+            all_cap_pos_ids.append(cap_padded_pos_ids)
+            # pad mask
+            cap_pad_mask = torch.cat(
+                [
+                    torch.zeros((cap_ori_len,), dtype=torch.bool, device=device),
+                    torch.ones((cap_padding_len,), dtype=torch.bool, device=device),
+                ],
+                dim=0,
+            )
+            all_cap_pad_mask.append(
+                cap_pad_mask if cap_padding_len > 0 else torch.zeros((cap_ori_len,), dtype=torch.bool, device=device)
+            )
+
+            if cap_feat_2 is not None:
+                if cap_feat is not None:
+                    all_cap_feats_out.append(cap_feat)
+                cap_padded_feat = torch.cat([cap_feat_2, (cap_feat_2[-1:] * 0).repeat(cap_padding_len, 1)], dim=0)
+                all_cap_feats_2_out.append(cap_padded_feat)
+            else:
+                cap_padded_feat = torch.cat([cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], dim=0)
+                all_cap_feats_out.append(cap_padded_feat)
+
+            ### Process Image
+            if image_ref is not None:
+                image = torch.cat([image, image_ref], dim=1)
+
+            C, F, H, W = image.size()
+            all_image_size.append((F, H, W))
+            F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
+
+            image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW)
+            # "c f pf h ph w pw -> (f h w) (pf ph pw c)"
+            image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(F_tokens * H_tokens * W_tokens, pF * pH * pW * C)
+
+            image_ori_len = len(image)
+            image_padding_len = (-image_ori_len) % SEQ_MULTI_OF
+
+            image_ori_pos_ids = self.create_coordinate_grid(
+                size=(F_tokens, H_tokens, W_tokens),
+                start=(cap_ori_len + cap_padding_len + 1, 0, 0),
+                device=device,
+            ).flatten(0, 2)
+            image_padded_pos_ids = torch.cat(
+                [
+                    image_ori_pos_ids,
+                    self.create_coordinate_grid(size=(1, 1, 1), start=(0, 0, 0), device=device)
+                    .flatten(0, 2)
+                    .repeat(image_padding_len, 1),
+                ],
+                dim=0,
+            )
+            all_image_pos_ids.append(image_padded_pos_ids if image_padding_len > 0 else image_ori_pos_ids)
+            # pad mask
+            image_pad_mask = torch.cat(
+                [
+                    torch.zeros((image_ori_len,), dtype=torch.bool, device=device),
+                    torch.ones((image_padding_len,), dtype=torch.bool, device=device),
+                ],
+                dim=0,
+            )
+            all_image_pad_mask.append(
+                image_pad_mask
+                if image_padding_len > 0
+                else torch.zeros((image_ori_len,), dtype=torch.bool, device=device)
+            )
+            # padded feature
+            image_padded_feat = torch.cat(
+                [image, image[-1:].repeat(image_padding_len, 1)],
+                dim=0,
+            )
+            all_image_out.append(image_padded_feat if image_padding_len > 0 else image)
+
+        return (
+            all_image_out,
+            all_cap_feats_out,
+            all_image_size,
+            all_image_pos_ids,
+            all_cap_pos_ids,
+            all_image_pad_mask,
+            all_cap_pad_mask,
+            all_cap_feats_2_out,
+        )
+
+    def forward(
+        self,
+        x: List[torch.Tensor],
+        t,
+        cap_feats: List[torch.Tensor],
+        patch_size=2,
+        f_patch_size=1,
+        ref_x=None,
+        cap_feats_2=None,
+        return_dict: bool = True,
+    ):
+        assert patch_size in self.all_patch_size
+        assert f_patch_size in self.all_f_patch_size
+
+        bsz = len(x)
+        device = x[0].device
+        t = t * self.t_scale
+        t = self.t_embedder(t)
+        
+        (
+            x,
+            cap_feats,
+            x_size,
+            x_pos_ids,
+            cap_pos_ids,
+            x_inner_pad_mask,
+            cap_inner_pad_mask,
+            cap_feats_2,
+        ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size, ref_x, cap_feats_2)
+
+        # x embed & refine
+        x_item_seqlens = [len(_) for _ in x]
+        assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens)
+        x_max_item_seqlen = max(x_item_seqlens)
+
+        x = torch.cat(x, dim=0)
+        x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x)
+
+        # Match t_embedder output dtype to x for layerwise casting compatibility
+        adaln_input = t.type_as(x)
+        if self.alignment_padding_mode == LEARNED_PADDING:
+            x[torch.cat(x_inner_pad_mask)] = self.x_pad_token
+        else:
+            x[torch.cat(x_inner_pad_mask)] = 0.0
+        x = list(x.split(x_item_seqlens, dim=0))
+        x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split([len(_) for _ in x_pos_ids], dim=0))
+
+        x = pad_sequence(x, batch_first=True, padding_value=0.0)
+        x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0)
+        # Clarify the length matches to satisfy Dynamo due to "Symbolic Shape Inference" to avoid compilation errors
+        x_freqs_cis = x_freqs_cis[:, : x.shape[1]]
+
+        x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(x_item_seqlens):
+            x_attn_mask[i, :seq_len] = 1
+        if self.alignment_padding_mode == ZERO_MASKED_PADDING:
+            mask_out_alignment_padding(x_attn_mask, x_inner_pad_mask, [0] * bsz)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for layer in self.noise_refiner:
+                x = self._gradient_checkpointing_func(layer, x, x_attn_mask, x_freqs_cis, adaln_input)
+        else:
+            for layer in self.noise_refiner:
+                x = layer(x, x_attn_mask, x_freqs_cis, adaln_input)
+
+
+        if len(cap_feats) > 0:
+            # cap embed & refine
+            cap_item_seqlens = [len(_) for _ in cap_feats]
+            cap_feats = torch.cat(cap_feats, dim=0)
+            cap_feats = self.cap_embedder(cap_feats)
+
+            if len(cap_feats_2) > 0:
+                cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0))
+                assert len(cap_feats) == len(cap_feats_2)
+                assert cap_feats[0].ndim == 2
+                assert cap_feats_2[0].ndim == 2
+                cap_feats = [torch.cat([i, j], dim=0) for i, j in zip(cap_feats, cap_feats_2)]
+                cap_item_seqlens = [len(_) for _ in cap_feats]
+                cap_feats = torch.cat(cap_feats, dim=0)
+        else:
+            assert len(cap_feats_2) > 0
+            cap_item_seqlens = [len(_) for _ in cap_feats_2]
+            cap_feats = torch.cat(cap_feats_2, dim=0)
+        
+        cap_max_item_seqlen = max(cap_item_seqlens)
+        if self.alignment_padding_mode == LEARNED_PADDING:
+            cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token
+        else:
+            cap_feats[torch.cat(cap_inner_pad_mask)] = 0.0
+
+        cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0))
+        cap_freqs_cis = list(
+            self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split([len(_) for _ in cap_pos_ids], dim=0)
+        )
+
+        cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0)
+        cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0)
+        # Clarify the length matches to satisfy Dynamo due to "Symbolic Shape Inference" to avoid compilation errors
+        cap_freqs_cis = cap_freqs_cis[:, : cap_feats.shape[1]]
+
+        cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(cap_item_seqlens):
+            cap_attn_mask[i, :seq_len] = 1
+        if self.alignment_padding_mode == ZERO_MASKED_PADDING:
+            mask_out_alignment_padding(cap_attn_mask, cap_inner_pad_mask, [0] * bsz)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for layer in self.context_refiner:
+                cap_feats = self._gradient_checkpointing_func(layer, cap_feats, cap_attn_mask, cap_freqs_cis)
+        else:
+            for layer in self.context_refiner:
+                cap_feats = layer(cap_feats, cap_attn_mask, cap_freqs_cis)
+
+        # unified
+        unified = []
+        unified_freqs_cis = []
+        for i in range(bsz):
+            x_len = x_item_seqlens[i]
+            cap_len = cap_item_seqlens[i]
+            unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]]))
+            unified_freqs_cis.append(torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]]))
+        unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)]
+        assert unified_item_seqlens == [len(_) for _ in unified]
+        unified_max_item_seqlen = max(unified_item_seqlens)
+
+        unified = pad_sequence(unified, batch_first=True, padding_value=0.0)
+        unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0)
+        unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(unified_item_seqlens):
+            unified_attn_mask[i, :seq_len] = 1
+        if self.alignment_padding_mode == ZERO_MASKED_PADDING:
+            # Unified layout is [image, caption] for every batch item.
+            mask_out_alignment_padding(unified_attn_mask, x_inner_pad_mask, [0] * bsz)
+            mask_out_alignment_padding(
+                unified_attn_mask, cap_inner_pad_mask, x_item_seqlens
+            )
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for layer in self.layers:
+                unified = self._gradient_checkpointing_func(
+                    layer, unified, unified_attn_mask, unified_freqs_cis, adaln_input
+                )
+        else:
+            for layer in self.layers:
+                unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input)
+
+        unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input)
+        unified = list(unified.unbind(dim=0))
+        x = self.unpatchify(unified, x_size, patch_size, f_patch_size)
+
+        if not return_dict:
+            return (x,)
+
+        return Transformer2DModelOutput(sample=x)

code/quant/load_int8.py

@@ -0,0 +1,172 @@
+"""Load a streamed INT8 Ming MLLM checkpoint onto a meta-initialized model.
+
+``model`` must already exist with parameters on ``meta`` (for example under
+``accelerate.init_empty_weights()``). Quantized modules listed in
+``int8_manifest.json`` are swapped from ``nn.Linear`` to ``Int8Linear.shell``
+before the shards are assigned in.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from torch import nn
+
+try:  # imported as the `quant` package (modeling_bailingmm2.py)
+    from .int8_linear import Int8Linear
+except ImportError:  # run from inside quant/ (CLI, tests)
+    from int8_linear import Int8Linear
+
+MANIFEST_NAME = "int8_manifest.json"
+INDEX_NAME = "model.safetensors.index.json"
+
+
+def load_int8_mllm_(model: nn.Module, int8_dir, device) -> dict:
+    """Swap quantize-rule linears for INT8 shells and assign shard tensors.
+
+    Returns ``{"modules_swapped", "tensors_loaded", "bytes_loaded"}``.
+    Raises ``RuntimeError`` on a bad manifest, a module that is not an
+    ``nn.Linear``, an unexpected checkpoint key, or any parameter / persistent
+    buffer still on ``meta``. Non-persistent buffers (rotary ``inv_freq``) may
+    stay on CPU; the caller moves the model afterwards.
+    """
+    int8_dir = Path(int8_dir)
+    dev = torch.device(device) if not isinstance(device, torch.device) else device
+    manifest_path = int8_dir / MANIFEST_NAME
+    if not manifest_path.is_file():
+        raise RuntimeError(f"missing int8 manifest: {manifest_path}")
+    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    if manifest.get("format") != "ming-int8-wo-v1":
+        raise RuntimeError(
+            f"unsupported int8 manifest format: {manifest.get('format')!r} ({manifest_path})"
+        )
+    module_names = manifest.get("quantized_modules")
+    if not isinstance(module_names, list) or not all(isinstance(n, str) for n in module_names):
+        raise RuntimeError(f"{manifest_path} quantized_modules is not a list of strings")
+
+    swapped = _swap_linears(model, module_names)
+
+    index_path = int8_dir / INDEX_NAME
+    if not index_path.is_file():
+        raise RuntimeError(f"missing index: {index_path}")
+    index = json.loads(index_path.read_text(encoding="utf-8"))
+    weight_map = index.get("weight_map")
+    if not isinstance(weight_map, dict) or not weight_map:
+        raise RuntimeError(f"{index_path} has no weight_map")
+
+    shard_names: list[str] = []
+    seen: set[str] = set()
+    for shard in weight_map.values():
+        if shard not in seen:
+            seen.add(shard)
+            shard_names.append(shard)
+
+    tensors_loaded = 0
+    bytes_loaded = 0
+    unexpected: list[str] = []
+    for shard in shard_names:
+        rel = Path(shard)
+        if rel.is_absolute() or ".." in rel.parts:
+            raise RuntimeError(f"unsafe shard path in index: {shard}")
+        path = int8_dir / rel
+        if not path.is_file():
+            raise RuntimeError(f"missing shard: {path}")
+        sd = load_file(str(path), device=str(dev))
+        for tensor in sd.values():
+            tensors_loaded += 1
+            bytes_loaded += tensor.numel() * tensor.element_size()
+        incompatible = model.load_state_dict(sd, strict=False, assign=True)
+        unexpected.extend(incompatible.unexpected_keys)
+        del sd
+
+    if unexpected:
+        listed = "\n".join(f"  {key}" for key in unexpected)
+        raise RuntimeError(
+            f"unexpected keys in checkpoint (not present on the model):\n{listed}"
+        )
+
+    _assert_loaded(model, module_names, dev)
+    return {
+        "modules_swapped": swapped,
+        "tensors_loaded": tensors_loaded,
+        "bytes_loaded": bytes_loaded,
+    }
+
+
+def _swap_linears(model: nn.Module, module_names: list[str]) -> int:
+    for name in module_names:
+        try:
+            linear = model.get_submodule(name)
+        except AttributeError as exc:
+            raise RuntimeError(f"manifest module not found on model: {name}") from exc
+        if not isinstance(linear, nn.Linear):
+            raise RuntimeError(
+                f"{name} is {type(linear).__name__}, expected nn.Linear "
+                "(refusing to swap a router or other non-linear)"
+            )
+        parent_name, _, leaf = name.rpartition(".")
+        if not leaf:
+            raise RuntimeError(f"cannot place shell for {name}")
+        parent = model.get_submodule(parent_name) if parent_name else model
+        has_bias = linear.bias is not None
+        bias_dtype = linear.bias.dtype if has_bias else torch.float32
+        shell = Int8Linear.shell(
+            in_features=linear.in_features,
+            out_features=linear.out_features,
+            bias=has_bias,
+            bias_dtype=bias_dtype,
+            device="meta",
+        )
+        setattr(parent, leaf, shell)
+    return len(module_names)
+
+
+def _assert_loaded(model: nn.Module, module_names: list[str], dev: torch.device) -> None:
+    offenders: list[str] = []
+    for name, param in model.named_parameters(remove_duplicate=False):
+        if param is not None and param.device.type == "meta":
+            offenders.append(f"parameter {name} dtype={param.dtype} device={param.device}")
+    for mod_name, mod in model.named_modules():
+        nonpersist = getattr(mod, "_non_persistent_buffers_set", set())
+        for buf_name, buf in mod._buffers.items():
+            if buf is None:
+                continue
+            full = f"{mod_name}.{buf_name}" if mod_name else buf_name
+            if buf.device.type != "meta":
+                # Non-persistent buffers (rotary inv_freq) are not in the
+                # checkpoint. accelerate leaves them on CPU; that is not an error.
+                continue
+            if buf_name in nonpersist:
+                offenders.append(
+                    f"non-persistent buffer {full} dtype={buf.dtype} device={buf.device}"
+                )
+            else:
+                offenders.append(f"buffer {full} dtype={buf.dtype} device={buf.device}")
+    if offenders:
+        listed = "\n".join(f"  {line}" for line in offenders)
+        raise RuntimeError(f"tensors still on meta after load:\n{listed}")
+
+    for name in module_names:
+        mod = model.get_submodule(name)
+        if not isinstance(mod, Int8Linear):
+            raise RuntimeError(f"{name} was not swapped to Int8Linear")
+        if mod.weight is None or mod.weight.dtype != torch.int8:
+            raise RuntimeError(f"{name}.weight is not int8 after load")
+        if mod.scale is None or mod.scale.dtype != torch.float32:
+            raise RuntimeError(f"{name}.scale is not float32 after load")
+        if mod.weight.device.type == "meta" or mod.scale.device.type == "meta":
+            raise RuntimeError(f"{name} still has meta tensors after load")
+        if mod.weight.device != dev or mod.scale.device != dev:
+            raise RuntimeError(
+                f"{name} loaded on weight={mod.weight.device} scale={mod.scale.device}, "
+                f"expected {dev}"
+            )
+        if tuple(mod.scale.shape) != (mod.out_features,):
+            raise RuntimeError(
+                f"{name}.scale shape {tuple(mod.scale.shape)} != ({mod.out_features},)"
+            )
+        if mod.bias is not None and mod.bias.device != dev:
+            raise RuntimeError(f"{name}.bias is on {mod.bias.device}, expected {dev}")

code/quant/test_int8.py

@@ -0,0 +1,665 @@
+"""CPU tests for weight-only INT8 Ming MLLM quantize + load.
+
+Run: HIP_VISIBLE_DEVICES=-1 python test_int8.py
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+import tempfile
+import traceback
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+import quantize_stream
+from int8_linear import Int8Linear, is_quantizable, quantize_weight
+from load_int8 import load_int8_mllm_
+
+# Tiny stand-in for Ming's MLLM names. Not the real model.
+HIDDEN = 32
+INTER = 48
+VOCAB = 64
+N_EXPERTS = 2
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        var = x.float().pow(2).mean(dim=-1, keepdim=True)
+        y = x * torch.rsqrt(var + self.eps)
+        return (y * self.weight).to(dtype=x.dtype)
+
+
+class Attention(nn.Module):
+    def __init__(self, hidden: int):
+        super().__init__()
+        self.hidden = hidden
+        self.query_key_value = nn.Linear(hidden, hidden * 3, bias=True)
+        self.dense = nn.Linear(hidden, hidden, bias=False)
+        self.q_norm = RMSNorm(hidden)
+        self.k_norm = RMSNorm(hidden)
+        # Non-persistent, like BailingMoeV2RotaryEmbedding.inv_freq.
+        self.register_buffer(
+            "inv_freq", torch.arange(hidden // 2, dtype=torch.float32), persistent=False
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        qkv = self.query_key_value(x)
+        h = self.hidden
+        q = self.q_norm(qkv[..., :h])
+        k = self.k_norm(qkv[..., h : 2 * h])
+        v = qkv[..., 2 * h :]
+        return self.dense(q + k + v)
+
+
+class DenseMLP(nn.Module):
+    def __init__(self, hidden: int, inter: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden, inter, bias=False)
+        self.up_proj = nn.Linear(hidden, inter, bias=True)
+        self.down_proj = nn.Linear(inter, hidden, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Expert(nn.Module):
+    def __init__(self, hidden: int, inter: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden, inter, bias=False)
+        self.up_proj = nn.Linear(hidden, inter, bias=True)
+        self.down_proj = nn.Linear(inter, hidden, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Router(nn.Module):
+    """Not an nn.Linear. Leaf name is gate / image_gate / audio_gate."""
+
+    def __init__(self, hidden: int, n_experts: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(n_experts, hidden))
+        self.expert_bias = nn.Parameter(torch.zeros(n_experts), requires_grad=False)
+        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.weight, self.expert_bias)
+
+
+class MoeMLP(nn.Module):
+    def __init__(self, hidden: int, inter: int, n_experts: int):
+        super().__init__()
+        self.gate = Router(hidden, n_experts)
+        self.image_gate = Router(hidden, n_experts)
+        self.audio_gate = Router(hidden, n_experts)
+        self.experts = nn.ModuleList(Expert(hidden, inter) for _ in range(n_experts))
+        self.shared_experts = Expert(hidden, inter)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        scores = self.gate(x) + self.image_gate(x) + self.audio_gate(x)
+        weights = torch.softmax(scores, dim=-1)
+        mixed = self.shared_experts(x)
+        for i, expert in enumerate(self.experts):
+            mixed = mixed + expert(x) * weights[..., i : i + 1]
+        return mixed
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, hidden: int, mlp: nn.Module):
+        super().__init__()
+        self.input_layernorm = RMSNorm(hidden)
+        self.post_attention_layernorm = RMSNorm(hidden)
+        self.attention = Attention(hidden)
+        self.mlp = mlp
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attention(self.input_layernorm(x))
+        x = x + self.mlp(self.post_attention_layernorm(x))
+        return x
+
+
+class TinyMing(nn.Module):
+    """Names match the real checkpoint: model.model.layers.*, model.lm_head, vision.*."""
+
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Module()
+        self.model.model = nn.Module()
+        self.model.model.word_embeddings = nn.Embedding(VOCAB, HIDDEN)
+        self.model.model.layers = nn.ModuleList(
+            [
+                DecoderLayer(HIDDEN, DenseMLP(HIDDEN, INTER)),
+                DecoderLayer(HIDDEN, MoeMLP(HIDDEN, INTER, N_EXPERTS)),
+            ]
+        )
+        self.model.model.norm = RMSNorm(HIDDEN)
+        self.model.lm_head = nn.Linear(HIDDEN, VOCAB, bias=False)
+        block = nn.Module()
+        block.attn = nn.Module()
+        block.attn.qkv = nn.Linear(HIDDEN, HIDDEN, bias=False)
+        self.vision = nn.Module()
+        self.vision.blocks = nn.ModuleList([block])
+        self.linear_proj = nn.ModuleList([nn.Linear(HIDDEN, HIDDEN, bias=True)])
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        h = self.model.model.word_embeddings(input_ids)
+        for layer in self.model.model.layers:
+            h = layer(h)
+        h = self.model.model.norm(h)
+        return self.model.lm_head(h)
+
+
+# Modules the rule must select for TinyMing. Hardcoded — not derived from is_quantizable.
+EXPECTED_QUANT_MODULES = [
+    "model.model.layers.0.attention.dense",
+    "model.model.layers.0.attention.query_key_value",
+    "model.model.layers.0.mlp.down_proj",
+    "model.model.layers.0.mlp.gate_proj",
+    "model.model.layers.0.mlp.up_proj",
+    "model.model.layers.1.attention.dense",
+    "model.model.layers.1.attention.query_key_value",
+    "model.model.layers.1.mlp.experts.0.down_proj",
+    "model.model.layers.1.mlp.experts.0.gate_proj",
+    "model.model.layers.1.mlp.experts.0.up_proj",
+    "model.model.layers.1.mlp.experts.1.down_proj",
+    "model.model.layers.1.mlp.experts.1.gate_proj",
+    "model.model.layers.1.mlp.experts.1.up_proj",
+    "model.model.layers.1.mlp.shared_experts.down_proj",
+    "model.model.layers.1.mlp.shared_experts.gate_proj",
+    "model.model.layers.1.mlp.shared_experts.up_proj",
+]
+
+MUST_NOT_QUANTIZE = [
+    "model.model.layers.1.mlp.gate",
+    "model.model.layers.1.mlp.image_gate",
+    "model.model.layers.1.mlp.audio_gate",
+    "model.model.word_embeddings",
+    "model.model.norm",
+    "model.lm_head",
+    "vision.blocks.0.attn.qkv",
+    "linear_proj.0",
+    "model.model.layers.0.attention.q_norm",
+    "model.model.layers.0.input_layernorm",
+]
+
+
+def _move_parameters_to_meta(model: nn.Module) -> nn.Module:
+    """Parameters → meta, buffers stay where they are (CPU). Matches accelerate include_buffers=False."""
+    for mod in model.modules():
+        for name, param in list(mod._parameters.items()):
+            if param is None:
+                continue
+            mod._parameters[name] = nn.Parameter(
+                param.detach().to(device="meta"),
+                requires_grad=param.requires_grad,
+            )
+    return model
+
+
+def _save_bf16_checkpoint(model: nn.Module, src: Path) -> None:
+    src.mkdir(parents=True, exist_ok=True)
+    sd = {k: v.detach().contiguous() for k, v in model.state_dict().items()}
+    if not sd:
+        raise AssertionError("empty state_dict")
+    for tensor in sd.values():
+        if tensor.is_floating_point():
+            assert tensor.dtype == torch.bfloat16, tensor.dtype
+    keys = list(sd)
+    mid = max(1, len(keys) // 2)
+    shards = {
+        "bf16-00001.safetensors": {k: sd[k] for k in keys[:mid]},
+        "bf16-00002.safetensors": {k: sd[k] for k in keys[mid:]},
+    }
+    weight_map = {}
+    total = 0
+    for filename, tensors in shards.items():
+        save_file(tensors, str(src / filename))
+        for name, tensor in tensors.items():
+            weight_map[name] = filename
+            total += tensor.numel() * tensor.element_size()
+    index = {"metadata": {"total_size": total}, "weight_map": weight_map}
+    (src / "model.safetensors.index.json").write_text(
+        json.dumps(index, indent=2) + "\n", encoding="utf-8"
+    )
+    (src / "config.json").write_bytes(b'{"model_type":"tiny-ming","hidden":32}\n')
+    extra = src / "extra"
+    extra.mkdir()
+    (extra / "chat_template.jinja").write_text("{{ messages }}\n", encoding="utf-8")
+
+
+def _load_all(folder: Path) -> dict[str, torch.Tensor]:
+    index = json.loads((folder / "model.safetensors.index.json").read_text(encoding="utf-8"))
+    order: list[str] = []
+    seen: set[str] = set()
+    for shard in index["weight_map"].values():
+        if shard not in seen:
+            seen.add(shard)
+            order.append(shard)
+    sd: dict[str, torch.Tensor] = {}
+    for shard in order:
+        sd.update(load_file(str(folder / shard)))
+    return sd
+
+
+def _apply_int8_(model: nn.Module) -> None:
+    names = []
+    for name, mod in model.named_modules():
+        if isinstance(mod, nn.Linear) and is_quantizable(
+            f"{name}.weight", tuple(mod.weight.shape)
+        ):
+            names.append(name)
+    for name in names:
+        parent_name, _, leaf = name.rpartition(".")
+        parent = model.get_submodule(parent_name) if parent_name else model
+        setattr(parent, leaf, Int8Linear.from_linear(getattr(parent, leaf)))
+
+
+def _assert_no_meta(model: nn.Module) -> None:
+    for name, param in model.named_parameters():
+        assert param.device.type != "meta", name
+    for mod_name, mod in model.named_modules():
+        for buf_name, buf in mod._buffers.items():
+            if buf is None:
+                continue
+            full = f"{mod_name}.{buf_name}" if mod_name else buf_name
+            assert buf.device.type != "meta", full
+
+
+def test_from_linear_roundtrip() -> None:
+    torch.manual_seed(0)
+    out_f, in_f = 5, 7
+    lin = nn.Linear(in_f, out_f, bias=True)
+    scales = torch.tensor([0.5, 0.25, 0.125, 2.0, 4.0], dtype=torch.float32)
+    q = torch.randint(-127, 128, (out_f, in_f), dtype=torch.int8)
+    q[:, 0] = 127
+    q[2, :] = 0  # all-zero row; must not NaN
+    weight = q.float() * scales[:, None]
+    with torch.no_grad():
+        lin.weight.copy_(weight)
+        lin.bias.copy_(torch.tensor([0.1, -0.2, 0.3, -0.4, 0.5]))
+    mod = Int8Linear.from_linear(lin)
+    deq = mod.weight.float() * mod.scale[:, None]
+    for row in range(out_f):
+        if row == 2:
+            assert torch.equal(mod.weight[row], torch.zeros(in_f, dtype=torch.int8))
+            assert float(mod.scale[row]) == 1.0
+            assert torch.equal(deq[row], torch.zeros(in_f))
+        else:
+            assert torch.equal(deq[row], weight[row]), (deq[row] - weight[row]).abs().max().item()
+    assert mod.bias is not None and torch.equal(mod.bias, lin.bias)
+    assert mod.bias.dtype == lin.bias.dtype
+    assert torch.isfinite(mod.scale).all()
+
+    # Random weights: per-element error stays within half a bin (+ float slack).
+    lin_r = nn.Linear(13, 9, bias=False)
+    mod_r = Int8Linear.from_linear(lin_r)
+    w = lin_r.weight.detach().float()
+    deq_r = (mod_r.weight.double() * mod_r.scale.double()[:, None]).float()
+    err = (w.double() - deq_r.double()).abs()
+    half = mod_r.scale.double()[:, None] * 0.5
+    slip = (err - half).max().item()
+    assert slip <= 1e-4, slip
+    assert torch.isfinite(mod_r.scale).all()
+
+    # Entirely zero weight: finite forward, zero codes, scale 1.
+    lin_z = nn.Linear(4, 3, bias=True)
+    with torch.no_grad():
+        lin_z.weight.zero_()
+    mod_z = Int8Linear.from_linear(lin_z)
+    assert torch.equal(mod_z.weight, torch.zeros_like(mod_z.weight))
+    assert torch.equal(mod_z.scale, torch.ones(3))
+    y = mod_z(torch.randn(8, 4))
+    assert torch.isfinite(y).all()
+    assert torch.allclose(y, mod_z.bias.expand_as(y))
+
+    # Zero row contributes only its bias.
+    x = torch.randn(6, in_f)
+    y_mix = mod(x)
+    assert torch.isfinite(y_mix).all()
+    assert torch.allclose(y_mix[:, 2], mod.bias[2].expand(6))
+
+    # bf16 source linear: codes int8, scale fp32, bias stays bf16.
+    lin_b = nn.Linear(8, 4, bias=True).to(dtype=torch.bfloat16)
+    mod_b = Int8Linear.from_linear(lin_b)
+    assert mod_b.weight.dtype == torch.int8
+    assert mod_b.scale.dtype == torch.float32
+    assert mod_b.bias is not None and mod_b.bias.dtype == torch.bfloat16
+    w_b = lin_b.weight.detach().float()
+    deq_b = mod_b.weight.float() * mod_b.scale[:, None]
+    err_b = (w_b.double() - deq_b.double()).abs()
+    half_b = mod_b.scale.double()[:, None] * 0.5
+    assert (err_b - half_b).max().item() <= 1e-2, (err_b - half_b).max().item()
+
+
+def _assert_quant_dtypes(mod: Int8Linear, scale: torch.Tensor, weight: torch.Tensor, bias_dtype: torch.dtype) -> None:
+    assert mod.weight.dtype == torch.int8
+    assert mod.scale.dtype == torch.float32
+    assert torch.equal(mod.weight, weight)
+    assert torch.equal(mod.scale, scale)
+    assert mod.bias is not None and mod.bias.dtype == bias_dtype
+
+
+def test_dtype_cast_keeps_scale_fp32() -> None:
+    torch.manual_seed(1)
+    lin = nn.Linear(5, 3, bias=True)
+    fresh = Int8Linear.from_linear(lin)
+    scale = fresh.scale.detach().clone()
+    weight = fresh.weight.detach().clone()
+    bias = fresh.bias.detach().clone()
+    assert scale.dtype == torch.float32 and weight.dtype == torch.int8 and bias.dtype == torch.float32
+
+    # Each cast starts from fp32 so "bias follows the cast" is the single cast of the source bias.
+    mod = Int8Linear.from_linear(lin)
+    mod.bfloat16()
+    _assert_quant_dtypes(mod, scale, weight, torch.bfloat16)
+    assert torch.equal(mod.bias, bias.to(dtype=torch.bfloat16))
+
+    mod = Int8Linear.from_linear(lin)
+    mod.half()
+    _assert_quant_dtypes(mod, scale, weight, torch.float16)
+    assert torch.equal(mod.bias, bias.to(dtype=torch.float16))
+
+    mod = Int8Linear.from_linear(lin)
+    mod.to(torch.bfloat16)
+    _assert_quant_dtypes(mod, scale, weight, torch.bfloat16)
+    assert torch.equal(mod.bias, bias.to(dtype=torch.bfloat16))
+
+    mod = Int8Linear.from_linear(lin)
+    mod.to(dtype=torch.float16)
+    _assert_quant_dtypes(mod, scale, weight, torch.float16)
+    assert torch.equal(mod.bias, bias.to(dtype=torch.float16))
+
+    # A second cast applies to the bias's current dtype, not the original fp32 value.
+    mod = Int8Linear.from_linear(lin)
+    mod.to(torch.bfloat16)
+    mod.to(dtype=torch.float16)
+    _assert_quant_dtypes(mod, scale, weight, torch.float16)
+    assert torch.equal(mod.bias, bias.to(dtype=torch.bfloat16).to(dtype=torch.float16))
+
+    # What the caller actually does: parent.to(device=..., dtype=bf16).
+    parent = nn.Sequential(Int8Linear.from_linear(lin))
+    parent.to(device="cpu", dtype=torch.bfloat16)
+    _assert_quant_dtypes(parent[0], scale, weight, torch.bfloat16)
+    assert torch.equal(parent[0].bias, bias.to(dtype=torch.bfloat16))
+
+    shell = Int8Linear.shell(4, 3, bias=True, bias_dtype=torch.bfloat16, device="meta")
+    assert shell.weight.dtype == torch.int8 and shell.weight.device.type == "meta"
+    assert shell.scale.dtype == torch.float32 and shell.scale.device.type == "meta"
+    assert shell.bias is not None
+    assert shell.bias.dtype == torch.bfloat16 and shell.bias.device.type == "meta"
+    shell_nb = Int8Linear.shell(4, 3, bias=False, bias_dtype=torch.float32, device="meta")
+    assert shell_nb.bias is None
+
+
+def test_forward_matches_reference() -> None:
+    torch.manual_seed(2)
+    for bias in (True, False):
+        lin = nn.Linear(6, 4, bias=bias)
+        # Bias is passed through unchanged, so it has to already match x's dtype
+        # (the caller does model.to(dtype=...) before the prefill).
+        modules = [
+            (Int8Linear.from_linear(lin), torch.float32),
+            (Int8Linear.from_linear(lin).to(torch.bfloat16), torch.bfloat16),
+            (Int8Linear.from_linear(lin).to(dtype=torch.float16), torch.float16),
+        ]
+        for mod, dtype in modules:
+            if mod.bias is not None:
+                assert mod.bias.dtype == dtype
+            x = torch.randn(3, 5, 6, dtype=dtype)
+            ref_w = (mod.weight.float() * mod.scale[:, None]).to(dtype=x.dtype)
+            y = mod(x)
+            y_ref = F.linear(x, ref_w, mod.bias)
+            assert torch.equal(y, y_ref), (bias, dtype)
+
+
+def test_is_quantizable_rule() -> None:
+    false_cases = [
+        ("model.model.layers.1.mlp.gate.weight", (256, 2048)),
+        ("model.model.layers.1.mlp.image_gate.weight", (256, 2048)),
+        ("model.model.layers.1.mlp.audio_gate.weight", (256, 2048)),
+        ("model.model.layers.1.mlp.gate.expert_bias", (256,)),
+        ("model.lm_head.weight", (151936, 2048)),
+        ("model.model.word_embeddings.weight", (151936, 2048)),
+        ("vision.blocks.0.attn.qkv.weight", (3072, 1280)),
+        ("model.model.layers.0.input_layernorm.weight", (2048,)),
+        ("model.model.layers.0.post_attention_layernorm.weight", (2048,)),
+        ("model.model.layers.0.attention.q_norm.weight", (128,)),
+        ("model.model.layers.0.attention.k_norm.weight", (128,)),
+        ("model.model.norm.weight", (2048,)),
+        ("linear_proj.0.weight", (2048, 2048)),
+        ("model.model.layers.0.attention.query_key_value.bias", (3072,)),
+        ("model.model.layers.0.mlp.experts.0.gate_proj.bias", (512,)),
+        # Right leaf, wrong rank: not quantizable (the stream must reject it).
+        ("model.model.layers.0.attention.query_key_value.weight", (3072,)),
+        ("model.model.layers.0.mlp.gate_proj.weight", (1024, 2048, 1)),
+    ]
+    true_cases = [
+        ("model.model.layers.3.mlp.experts.3.gate_proj.weight", (512, 2048)),
+        ("model.model.layers.3.mlp.shared_experts.down_proj.weight", (2048, 512)),
+        ("model.model.layers.0.mlp.up_proj.weight", (512, 2048)),
+        ("layers.0.mlp.up_proj.weight", (512, 2048)),
+        ("model.model.layers.0.attention.query_key_value.weight", (3072, 2048)),
+        ("model.model.layers.0.attention.dense.weight", (2048, 2048)),
+        ("model.model.layers.0.mlp.gate_proj.weight", (512, 2048)),
+        ("model.model.layers.0.mlp.down_proj.weight", (2048, 512)),
+        ("model.model.layers.19.mlp.experts.255.up_proj.weight", (512, 2048)),
+    ]
+    for name, shape in false_cases:
+        assert is_quantizable(name, shape) is False, name
+    for name, shape in true_cases:
+        assert is_quantizable(name, shape) is True, name
+
+
+def _shard_groups(sd: dict[str, torch.Tensor]) -> set[str]:
+    """One copy-tensor, or one weight+scale pair, is one unsplittable group."""
+    names = set(sd)
+    groups: set[str] = set()
+    for name in names:
+        if name.endswith(".scale") and name[: -len(".scale")] + ".weight" in names:
+            groups.add(name[: -len(".scale")])
+        elif name.endswith(".weight") and name[: -len(".weight")] + ".scale" in names:
+            groups.add(name[: -len(".weight")])
+        else:
+            groups.add(name)
+    return groups
+
+
+def test_end_to_end_stream_and_load() -> None:
+    assert quantize_stream.MAX_SHARD_BYTES == 5 * 10**9
+    torch.manual_seed(3)
+    src_model = TinyMing().to(dtype=torch.bfloat16)
+    # Non-persistent rotary buffer is not part of the checkpoint.
+    assert "model.model.layers.0.attention.inv_freq" not in src_model.state_dict()
+
+    with tempfile.TemporaryDirectory(prefix="ming-int8-") as tmp:
+        root = Path(tmp)
+        src = root / "src"
+        dst = root / "dst"
+        _save_bf16_checkpoint(src_model, src)
+        limit = 2048
+        old = quantize_stream.MAX_SHARD_BYTES
+        quantize_stream.MAX_SHARD_BYTES = limit
+        try:
+            rc = quantize_stream.main([str(src), str(dst)])
+        finally:
+            quantize_stream.MAX_SHARD_BYTES = old
+        assert rc == 0, rc
+        assert quantize_stream.MAX_SHARD_BYTES == 5 * 10**9
+
+        # Sidecars copied verbatim; original index replaced.
+        assert (dst / "config.json").read_bytes() == (src / "config.json").read_bytes()
+        assert (dst / "extra" / "chat_template.jinja").read_bytes() == (
+            src / "extra" / "chat_template.jinja"
+        ).read_bytes()
+        assert not (dst / "bf16-00001.safetensors").exists()
+
+        manifest = json.loads((dst / "int8_manifest.json").read_text(encoding="utf-8"))
+        assert manifest["format"] == "ming-int8-wo-v1"
+        assert manifest["scheme"] == (
+            "weight-only int8, per-output-channel symmetric, fp32 scales"
+        )
+        assert manifest["quantized_modules"] == sorted(EXPECTED_QUANT_MODULES)
+        for banned in MUST_NOT_QUANTIZE:
+            assert banned not in manifest["quantized_modules"], banned
+
+        index = json.loads((dst / "model.safetensors.index.json").read_text(encoding="utf-8"))
+        assert index["metadata"]["total_size"] == manifest["total_size"]
+        measured = manifest["measured"]
+        assert measured["tensors_quantized"] == len(EXPECTED_QUANT_MODULES)
+        assert measured["bytes_in"] == manifest["source_total_size"]
+        assert measured["bytes_out"] == manifest["total_size"]
+        assert measured["bytes_out"] < measured["bytes_in"]
+        n_out_keys = measured["tensors_copied"] + 2 * measured["tensors_quantized"]
+        assert len(index["weight_map"]) == n_out_keys
+
+        src_sd = _load_all(src)
+        dst_sd = _load_all(dst)
+        assert manifest["source_total_size"] == sum(
+            t.numel() * t.element_size() for t in src_sd.values()
+        )
+        assert manifest["total_size"] == sum(t.numel() * t.element_size() for t in dst_sd.values())
+
+        shard_names = sorted({*index["weight_map"].values()})
+        assert len(shard_names) >= 2, shard_names
+        for shard in shard_names:
+            shard_sd = load_file(str(dst / shard))
+            total = sum(t.numel() * t.element_size() for t in shard_sd.values())
+            if total > limit:
+                assert len(_shard_groups(shard_sd)) == 1, (shard, total, list(shard_sd))
+
+        errors = []
+        for name, src_t in src_sd.items():
+            if is_quantizable(name, tuple(src_t.shape)):
+                q = dst_sd[name]
+                scale_key = name[: -len("weight")] + "scale"
+                scale = dst_sd[scale_key]
+                assert q.dtype == torch.int8, name
+                assert scale.dtype == torch.float32, scale_key
+                q_ref, scale_ref = quantize_weight(src_t)
+                assert torch.equal(q, q_ref), name
+                assert torch.equal(scale, scale_ref), scale_key
+                errors.append((name, quantize_stream._relative_frobenius(src_t, q, scale)))
+            else:
+                assert name in dst_sd, name
+                assert dst_sd[name].dtype == src_t.dtype, (name, dst_sd[name].dtype, src_t.dtype)
+                assert torch.equal(dst_sd[name], src_t), name
+        # Router weights stayed BF16 and byte-identical (the gate vs gate_proj trap).
+        router = "model.model.layers.1.mlp.gate.weight"
+        assert dst_sd[router].dtype == torch.bfloat16
+        assert torch.equal(dst_sd[router], src_sd[router])
+        for suffix in ("image_gate.weight", "audio_gate.weight", "gate.expert_bias"):
+            key = f"model.model.layers.1.mlp.{suffix}"
+            assert torch.equal(dst_sd[key], src_sd[key]), key
+
+        vals = [e for _, e in errors]
+        assert measured["max_relative_error"] == max(vals)
+        assert measured["mean_relative_error"] == sum(vals) / len(vals)
+        assert measured["worst_tensor"] in dict(errors)
+        assert measured["max_relative_error"] == dict(errors)[measured["worst_tensor"]]
+        assert 0.0 <= measured["mean_relative_error"] <= measured["p99_relative_error"]
+        assert measured["p99_relative_error"] <= measured["max_relative_error"]
+        assert measured["max_relative_error"] < 0.05, measured
+
+        # Eager quant of the same BF16 bytes.
+        eager = TinyMing().to(dtype=torch.bfloat16)
+        incompatible = eager.load_state_dict(src_sd, strict=True)
+        assert not incompatible.missing_keys and not incompatible.unexpected_keys
+        _apply_int8_(eager)
+
+        loaded = _move_parameters_to_meta(TinyMing())
+        for layer in loaded.model.model.layers:
+            assert layer.attention.inv_freq.device.type == "cpu"
+            assert layer.attention.query_key_value.weight.device.type == "meta"
+        report = load_int8_mllm_(loaded, dst, "cpu")
+        assert report["modules_swapped"] == len(EXPECTED_QUANT_MODULES)
+        assert report["tensors_loaded"] == len(dst_sd)
+        assert report["bytes_loaded"] == manifest["total_size"]
+        _assert_no_meta(loaded)
+        for layer in loaded.model.model.layers:
+            assert layer.attention.inv_freq.device.type == "cpu"
+            assert layer.attention.inv_freq.dtype == torch.float32
+        for name in EXPECTED_QUANT_MODULES:
+            mod = loaded.get_submodule(name)
+            assert isinstance(mod, Int8Linear), name
+            assert mod.weight.dtype == torch.int8
+            assert mod.scale.dtype == torch.float32
+
+        eager.eval()
+        loaded.eval()
+        ids = torch.randint(0, VOCAB, (2, 6))
+        with torch.no_grad():
+            y_eager = eager(ids)
+            y_loaded = loaded(ids)
+        assert y_eager.dtype == y_loaded.dtype
+        assert torch.equal(y_eager, y_loaded), (y_eager - y_loaded).abs().max().item()
+
+        # A second run into a non-empty safetensors dir must fail loudly.
+        print("  re-running into a non-empty dst (expect error on stderr)", flush=True)
+        rc_again = quantize_stream.main([str(src), str(dst)])
+        assert rc_again == 1
+
+
+def test_unknown_key_fails_loudly() -> None:
+    torch.manual_seed(4)
+    model = TinyMing().to(dtype=torch.bfloat16)
+    with tempfile.TemporaryDirectory(prefix="ming-int8-bad-") as tmp:
+        root = Path(tmp)
+        src = root / "src"
+        dst = root / "dst"
+        _save_bf16_checkpoint(model, src)
+        rc = quantize_stream.main([str(src), str(dst)])
+        assert rc == 0, rc
+        shard = next(dst.glob("*.safetensors"))
+        sd = load_file(str(shard))
+        sd["not.a.real.key"] = torch.zeros(4, dtype=torch.float32)
+        save_file(sd, str(shard))
+        loaded = _move_parameters_to_meta(TinyMing())
+        try:
+            load_int8_mllm_(loaded, dst, "cpu")
+        except RuntimeError as exc:
+            text = str(exc)
+            assert "unexpected" in text.lower(), text
+            assert "not.a.real.key" in text, text
+            print(f"  caught RuntimeError: {text.splitlines()[0]}")
+        else:
+            raise AssertionError("load_int8_mllm_ returned instead of failing on an unknown key")
+
+
+def main() -> int:
+    import safetensors
+
+    print(f"torch={torch.__version__} safetensors={safetensors.__version__}", flush=True)
+    tests = [
+        test_from_linear_roundtrip,
+        test_dtype_cast_keeps_scale_fp32,
+        test_forward_matches_reference,
+        test_is_quantizable_rule,
+        test_end_to_end_stream_and_load,
+        test_unknown_key_fails_loudly,
+    ]
+    failed = 0
+    for fn in tests:
+        try:
+            fn()
+        except Exception:
+            failed += 1
+            print(f"FAIL {fn.__name__}", flush=True)
+            traceback.print_exc()
+        else:
+            print(f"PASS {fn.__name__}", flush=True)
+    print(f"{len(tests) - failed} passed, {failed} failed", flush=True)
+    return 1 if failed else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

code/tests/test_infer_cli.py

@@ -0,0 +1,226 @@
+import json
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+import unittest
+from unittest.mock import patch
+
+from infer import parse_args, resolve_task_resolution
+
+
+REPOSITORY = Path(__file__).resolve().parents[1]
+INFER = REPOSITORY / "infer.py"
+
+
+class InferenceCliTest(unittest.TestCase):
+    def _model_directory(self, profile):
+        temporary = tempfile.TemporaryDirectory()
+        model_directory = Path(temporary.name)
+        (model_directory / "inference_profile.json").write_text(
+            json.dumps(profile), encoding="utf-8"
+        )
+        return temporary, model_directory
+
+    def test_cli_defaults_to_one_gpu(self):
+        with patch.object(
+            sys,
+            "argv",
+            ["infer.py", "--model", "checkpoint", "--task", "text-to-image"],
+        ):
+            args = parse_args()
+        self.assertEqual(args.device, "cuda:0")
+        self.assertEqual(args.device_map, "balanced")
+        self.assertEqual(args.num_gpus, 1)
+        self.assertIsNone(args.resolution)
+
+    def test_task_resolution_defaults_and_snapping(self):
+        self.assertEqual(resolve_task_resolution("text-to-image", None), 2048)
+        self.assertEqual(resolve_task_resolution("text-to-image", 1200), 1024)
+        self.assertEqual(resolve_task_resolution("text-to-image", 1800), 2048)
+        self.assertEqual(resolve_task_resolution("text-to-image", 1536), 1024)
+
+        self.assertEqual(resolve_task_resolution("image-edit", None), 1024)
+        self.assertEqual(resolve_task_resolution("image-edit", 512), 1024)
+        self.assertEqual(resolve_task_resolution("image-edit", 2048), 1024)
+
+        self.assertEqual(resolve_task_resolution("layer-decompose", None), 1024)
+        self.assertEqual(resolve_task_resolution("layer-decompose", 600), 512)
+        self.assertEqual(resolve_task_resolution("layer-decompose", 900), 1024)
+        self.assertEqual(resolve_task_resolution("layer-decompose", 768), 512)
+
+    def test_task_resolution_rejects_non_positive_values(self):
+        for value in (0, -1):
+            with self.subTest(value=value):
+                with self.assertRaisesRegex(ValueError, "positive integer"):
+                    resolve_task_resolution("text-to-image", value)
+
+    def test_validate_only_accepts_local_generation_checkpoint(self):
+        temporary, model_directory = self._model_directory(
+            {
+                "schema_version": 1,
+                "inference_profile": "generation_edit",
+                "alignment_padding_mode": "zero_masked",
+                "multi_frame_output": False,
+                "vae_input_channels": 4,
+                "vae_sample_mode": "argmax",
+            }
+        )
+        self.addCleanup(temporary.cleanup)
+        result = subprocess.run(
+            [
+                sys.executable,
+                str(INFER),
+                "--model",
+                str(model_directory),
+                "--task",
+                "text-to-image",
+                "--prompt",
+                "test",
+                "--validate-only",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        payload = json.loads(result.stdout)
+        self.assertEqual(payload["task"], "text-to-image")
+        self.assertEqual(payload["sampling"], {"steps": 12, "cfg": 1.0})
+        self.assertEqual(
+            payload["resolution"], {"requested": None, "effective": 2048}
+        )
+
+    def test_validate_only_accepts_long_literal_prompt(self):
+        temporary, model_directory = self._model_directory(
+            {
+                "schema_version": 1,
+                "inference_profile": "generation_edit",
+                "alignment_padding_mode": "zero_masked",
+                "multi_frame_output": False,
+                "vae_input_channels": 4,
+                "vae_sample_mode": "argmax",
+            }
+        )
+        self.addCleanup(temporary.cleanup)
+        long_prompt = "Create a detailed ocean research poster. " * 40
+        result = subprocess.run(
+            [
+                sys.executable,
+                str(INFER),
+                "--model",
+                str(model_directory),
+                "--task",
+                "text-to-image",
+                "--prompt",
+                long_prompt,
+                "--validate-only",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        payload = json.loads(result.stdout)
+        self.assertEqual(payload["task"], "text-to-image")
+        self.assertEqual(payload["sampling"], {"steps": 12, "cfg": 1.0})
+
+    def test_validate_only_uses_layer_defaults_and_accepts_overrides(self):
+        temporary, model_directory = self._model_directory(
+            {
+                "schema_version": 1,
+                "inference_profile": "layer_decompose",
+                "alignment_padding_mode": "learned",
+                "multi_frame_output": True,
+                "vae_input_channels": 4,
+                "vae_sample_mode": "argmax",
+            }
+        )
+        self.addCleanup(temporary.cleanup)
+        input_image = model_directory / "input.png"
+        input_image.write_bytes(b"validation-only")
+
+        default_result = subprocess.run(
+            [
+                sys.executable,
+                str(INFER),
+                "--model",
+                str(model_directory),
+                "--task",
+                "layer-decompose",
+                "--input-image",
+                str(input_image),
+                "--num-layers",
+                "4",
+                "--validate-only",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        self.assertEqual(
+            json.loads(default_result.stdout)["sampling"],
+            {"steps": 12, "cfg": 2.0},
+        )
+        self.assertEqual(
+            json.loads(default_result.stdout)["resolution"],
+            {"requested": None, "effective": 1024},
+        )
+
+        override_result = subprocess.run(
+            [
+                sys.executable,
+                str(INFER),
+                "--model",
+                str(model_directory),
+                "--task",
+                "layer-decompose",
+                "--input-image",
+                str(input_image),
+                "--steps",
+                "16",
+                "--cfg",
+                "1.25",
+                "--validate-only",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        self.assertEqual(
+            json.loads(override_result.stdout)["sampling"],
+            {"steps": 16, "cfg": 1.25},
+        )
+
+    def test_validate_only_rejects_wrong_checkpoint_family(self):
+        temporary, model_directory = self._model_directory(
+            {
+                "schema_version": 1,
+                "inference_profile": "layer_decompose",
+                "alignment_padding_mode": "learned",
+                "multi_frame_output": True,
+                "vae_input_channels": 4,
+                "vae_sample_mode": "argmax",
+            }
+        )
+        self.addCleanup(temporary.cleanup)
+        result = subprocess.run(
+            [
+                sys.executable,
+                str(INFER),
+                "--model",
+                str(model_directory),
+                "--task",
+                "text-to-image",
+                "--prompt",
+                "test",
+                "--validate-only",
+            ],
+            check=False,
+            capture_output=True,
+            text=True,
+        )
+        self.assertNotEqual(result.returncode, 0)
+        self.assertIn("generation_edit checkpoint", result.stderr)
+
+
+if __name__ == "__main__":
+    unittest.main()

code/tests/test_inference_profile.py

@@ -0,0 +1,257 @@
+import json
+import tempfile
+from pathlib import Path
+import unittest
+
+from inference_profile import (
+    InferenceProfile,
+    InferenceProfileError,
+    load_checkpoint_capabilities,
+    load_inference_profile,
+    resolve_model_directory,
+)
+
+
+GENERATION = {
+    "schema_version": 1,
+    "inference_profile": "generation_edit",
+    "alignment_padding_mode": "zero_masked",
+    "multi_frame_output": False,
+    "vae_input_channels": 4,
+    "vae_sample_mode": "argmax",
+}
+
+LAYER = {
+    "schema_version": 1,
+    "inference_profile": "layer_decompose",
+    "alignment_padding_mode": "learned",
+    "multi_frame_output": True,
+    "vae_input_channels": 4,
+    "vae_sample_mode": "argmax",
+}
+
+
+class InferenceProfileTest(unittest.TestCase):
+    def test_all_profile_fields_are_required(self):
+        for field in GENERATION:
+            raw = dict(GENERATION)
+            raw.pop(field)
+            with self.assertRaisesRegex(InferenceProfileError, "missing required fields"):
+                InferenceProfile.from_dict(raw)
+
+    def test_profile_rejects_unknown_fields(self):
+        raw = dict(GENERATION, inferred_from_directory_name=True)
+        with self.assertRaisesRegex(InferenceProfileError, "unsupported fields"):
+            InferenceProfile.from_dict(raw)
+
+    def test_generation_profile_task_matrix(self):
+        profile = InferenceProfile.from_dict(GENERATION)
+        self.assertEqual(profile.resolve_sampling_parameters().steps, 12)
+        self.assertEqual(profile.resolve_sampling_parameters().cfg, 1.0)
+        profile.validate_task("text-to-image", has_reference_image=False)
+        profile.validate_task("image-edit", has_reference_image=True)
+        with self.assertRaisesRegex(InferenceProfileError, "layer_decompose checkpoint"):
+            profile.validate_task("layer-decompose", has_reference_image=True, num_layers=4)
+
+    def test_layer_profile_task_matrix(self):
+        profile = InferenceProfile.from_dict(LAYER)
+        self.assertEqual(profile.resolve_sampling_parameters().steps, 12)
+        self.assertEqual(profile.resolve_sampling_parameters().cfg, 2.0)
+        profile.validate_task("layer-decompose", has_reference_image=True, num_layers=4)
+        with self.assertRaisesRegex(InferenceProfileError, "generation_edit checkpoint"):
+            profile.validate_task("image-edit", has_reference_image=True)
+
+    def test_generation_profile_requires_qwen_vae_contract(self):
+        with self.assertRaisesRegex(InferenceProfileError, "vae_input_channels=4"):
+            InferenceProfile.from_dict(dict(GENERATION, vae_input_channels=3))
+        with self.assertRaisesRegex(InferenceProfileError, "vae_sample_mode='argmax'"):
+            InferenceProfile.from_dict(dict(GENERATION, vae_sample_mode="sample"))
+
+    def test_layer_profile_requires_qwen_vae_contract(self):
+        with self.assertRaisesRegex(InferenceProfileError, "vae_sample_mode='argmax'"):
+            InferenceProfile.from_dict(dict(LAYER, vae_sample_mode="sample"))
+
+    def test_sampling_overrides_are_validated(self):
+        profile = InferenceProfile.from_dict(GENERATION)
+        resolved = profile.resolve_sampling_parameters(steps=18, cfg=1.5)
+        self.assertEqual((resolved.steps, resolved.cfg), (18, 1.5))
+        with self.assertRaisesRegex(InferenceProfileError, "steps"):
+            profile.resolve_sampling_parameters(steps=0)
+        with self.assertRaisesRegex(InferenceProfileError, "CFG"):
+            profile.resolve_sampling_parameters(cfg=float("inf"))
+
+    def test_profile_is_loaded_from_local_model_directory(self):
+        with self.subTest("local profile"):
+            import tempfile
+            from pathlib import Path
+
+            with tempfile.TemporaryDirectory() as directory:
+                model_directory = Path(directory)
+                (model_directory / "inference_profile.json").write_text(
+                    json.dumps(LAYER), encoding="utf-8"
+                )
+                self.assertEqual(
+                    load_inference_profile(model_directory).inference_profile,
+                    "layer_decompose",
+                )
+                self.assertEqual(
+                    resolve_model_directory(model_directory), model_directory.resolve()
+                )
+
+    def test_missing_profile_is_an_error(self):
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as directory:
+            with self.assertRaisesRegex(InferenceProfileError, "must contain"):
+                load_inference_profile(directory)
+
+    def test_missing_absolute_local_path_is_not_treated_as_hub_id(self):
+        import tempfile
+        from pathlib import Path
+
+        with tempfile.TemporaryDirectory() as directory:
+            with self.assertRaises(FileNotFoundError):
+                resolve_model_directory(Path(directory) / "missing")
+
+
+VAE_QWEN_4CH = {"_class_name": "AutoencoderKLQwenImage", "input_channels": 4}
+
+
+def build_component_package(
+    root: Path,
+    transformer_extra=None,
+    vae=None,
+    profile=None,
+) -> None:
+    (root / "transformer").mkdir(parents=True, exist_ok=True)
+    (root / "vae").mkdir(parents=True, exist_ok=True)
+    transformer = {"_class_name": "DiffusionTransformer", "dim": 4}
+    transformer.update(transformer_extra or {})
+    (root / "transformer" / "config.json").write_text(json.dumps(transformer))
+    (root / "vae" / "config.json").write_text(json.dumps(vae or VAE_QWEN_4CH))
+    if profile is not None:
+        (root / "inference_profile.json").write_text(json.dumps(profile))
+
+
+class ComponentCapabilityTest(unittest.TestCase):
+    def capabilities(self, **kwargs):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            build_component_package(root, **kwargs)
+            return load_checkpoint_capabilities(root)
+
+    def test_generation_pair_without_profile(self):
+        profile = self.capabilities(
+            transformer_extra={
+                "alignment_padding_mode": "zero_masked",
+                "multi_frame_output": False,
+            }
+        )
+        self.assertEqual(profile.inference_profile, "generation_edit")
+        self.assertEqual(profile.vae_input_channels, 4)
+        self.assertEqual(profile.vae_sample_mode, "argmax")
+        self.assertEqual(profile.resolve_sampling_parameters().steps, 12)
+        self.assertEqual(profile.resolve_sampling_parameters().cfg, 1.0)
+
+    def test_layer_pair_without_profile(self):
+        profile = self.capabilities(
+            transformer_extra={
+                "alignment_padding_mode": "learned",
+                "multi_frame_output": True,
+            }
+        )
+        self.assertEqual(profile.inference_profile, "layer_decompose")
+        self.assertEqual(profile.resolve_sampling_parameters().cfg, 2.0)
+
+    def test_partial_pair_is_rejected(self):
+        for field in ("alignment_padding_mode", "multi_frame_output"):
+            with self.subTest(only=field):
+                pair = {
+                    "alignment_padding_mode": "zero_masked",
+                    "multi_frame_output": False,
+                }
+                with self.assertRaisesRegex(InferenceProfileError, "declared together"):
+                    self.capabilities(transformer_extra={field: pair[field]})
+
+    def test_every_invalid_pair_is_rejected(self):
+        invalid = [
+            {"alignment_padding_mode": "zero_masked", "multi_frame_output": True},
+            {"alignment_padding_mode": "learned", "multi_frame_output": False},
+            {"alignment_padding_mode": "zero", "multi_frame_output": False},
+            {"alignment_padding_mode": "learned", "multi_frame_output": 1},
+            {"alignment_padding_mode": None, "multi_frame_output": False},
+        ]
+        for pair in invalid:
+            with self.subTest(pair=pair):
+                with self.assertRaises(InferenceProfileError):
+                    self.capabilities(transformer_extra=pair)
+
+    def test_vae_class_must_be_qwen(self):
+        with self.assertRaisesRegex(InferenceProfileError, "AutoencoderKLQwenImage"):
+            self.capabilities(
+                transformer_extra={
+                    "alignment_padding_mode": "zero_masked",
+                    "multi_frame_output": False,
+                },
+                vae={"_class_name": "AutoencoderKL", "in_channels": 4},
+            )
+
+    def test_vae_channel_fields_must_agree_and_be_four(self):
+        base = {"alignment_padding_mode": "zero_masked", "multi_frame_output": False}
+        with self.assertRaisesRegex(InferenceProfileError, "disagree"):
+            self.capabilities(
+                transformer_extra=base,
+                vae=dict(VAE_QWEN_4CH, in_channels=16),
+            )
+        with self.assertRaisesRegex(InferenceProfileError, "4-channel"):
+            self.capabilities(
+                transformer_extra=base,
+                vae={"_class_name": "AutoencoderKLQwenImage", "input_channels": 3},
+            )
+        with self.assertRaisesRegex(InferenceProfileError, "must declare"):
+            self.capabilities(
+                transformer_extra=base,
+                vae={"_class_name": "AutoencoderKLQwenImage"},
+            )
+        # in_channels alone is accepted when it agrees by itself.
+        profile = self.capabilities(
+            transformer_extra=base,
+            vae={"_class_name": "AutoencoderKLQwenImage", "in_channels": 4},
+        )
+        self.assertEqual(profile.vae_input_channels, 4)
+
+    def test_legacy_profile_fallback(self):
+        profile = self.capabilities(profile=LAYER)
+        self.assertEqual(profile.inference_profile, "layer_decompose")
+        self.assertEqual(profile.alignment_padding_mode, "learned")
+
+    def test_legacy_fallback_requires_profile_file(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            build_component_package(root)
+            with self.assertRaisesRegex(InferenceProfileError, "must contain"):
+                load_checkpoint_capabilities(root)
+
+    def test_component_fields_and_agreeing_legacy_profile(self):
+        profile = self.capabilities(
+            transformer_extra={
+                "alignment_padding_mode": "zero_masked",
+                "multi_frame_output": False,
+            },
+            profile=GENERATION,
+        )
+        self.assertEqual(profile.inference_profile, "generation_edit")
+
+    def test_component_fields_disagreeing_with_legacy_profile(self):
+        with self.assertRaisesRegex(InferenceProfileError, "disagree"):
+            self.capabilities(
+                transformer_extra={
+                    "alignment_padding_mode": "learned",
+                    "multi_frame_output": True,
+                },
+                profile=GENERATION,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

code/tests/test_inference_smoke.py

@@ -0,0 +1,206 @@
+"""Optional end-to-end GPU smoke tests for the maintained demo cases.
+
+Required environment variables:
+
+- MING_GENERATION_MODEL
+- MING_LAYER_MODEL
+
+The model variables may be local directories or Hugging Face Hub IDs.
+
+Optional environment variables:
+
+- MING_SMOKE_OUTPUT_DIR: keep outputs under this directory (created when
+  missing) instead of a temporary directory.
+- MING_SMOKE_MODE: "short" (default) runs the two-step probes only;
+  "full" additionally runs the default-parameter acceptance cases.
+- MING_SMOKE_DEVICE, MING_SMOKE_DEVICE_MAP, MING_SMOKE_NUM_GPUS,
+  MING_SMOKE_DTYPE and MING_SMOKE_ATTN_IMPLEMENTATION are forwarded to
+  infer.py. The smoke defaults to the validated single-GPU layout in BF16
+  and FlashAttention 2; set MING_SMOKE_NUM_GPUS (e.g. 8) for a multi-GPU
+  run.
+
+Text-to-image and layer decomposition use the same assets displayed in the
+public README. The image-edit compatibility probe keeps its established input
+and prompt.
+"""
+
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+import unittest
+
+
+REPOSITORY = Path(__file__).resolve().parents[1]
+INFER = REPOSITORY / "infer.py"
+T2I_PROMPT = REPOSITORY / "assets" / "t2i_four_seasons_cabin_prompt.json"
+IMAGE_EDIT_INPUT = REPOSITORY / "tests" / "assets" / "smoke_input.png"
+IMAGE_EDIT_PROMPT = "Change the background to blue"
+LAYER_INPUT = REPOSITORY / "assets" / "layer_samples" / "card_making_input.png"
+LAYER_PROMPT = REPOSITORY / "assets" / "layer_samples" / "card_making_prompt.txt"
+LAYER_COUNT = 6
+
+
+class SmokeAssetContractTest(unittest.TestCase):
+    def test_smoke_assets_are_present_and_well_formed(self):
+        for path in (T2I_PROMPT, IMAGE_EDIT_INPUT, LAYER_INPUT, LAYER_PROMPT):
+            self.assertTrue(path.is_file(), path)
+
+        with T2I_PROMPT.open(encoding="utf-8") as handle:
+            structured_prompt = json.load(handle)
+        self.assertEqual(set(structured_prompt), {"canvas_settings", "layers"})
+
+        layer_prompt = LAYER_PROMPT.read_text(encoding="utf-8")
+        self.assertIn(f"Number of layers: {LAYER_COUNT}", layer_prompt)
+
+
+def _smoke_output_root():
+    value = os.environ.get("MING_SMOKE_OUTPUT_DIR")
+    return Path(value).expanduser().resolve() if value else None
+
+
+@unittest.skipUnless(
+    os.environ.get("MING_GENERATION_MODEL")
+    and os.environ.get("MING_LAYER_MODEL"),
+    "set MING_GENERATION_MODEL and MING_LAYER_MODEL",
+)
+class InferenceSmokeTest(unittest.TestCase):
+    def _run(self, *arguments, steps=2, resolution=None):
+        command = [
+            sys.executable,
+            str(INFER),
+            *map(str, arguments),
+            "--seed",
+            "42",
+            "--device",
+            os.environ.get("MING_SMOKE_DEVICE", "cuda:0"),
+            "--device-map",
+            os.environ.get("MING_SMOKE_DEVICE_MAP", "balanced"),
+            "--num-gpus",
+            os.environ.get("MING_SMOKE_NUM_GPUS", "1"),
+            "--dtype",
+            os.environ.get("MING_SMOKE_DTYPE", "bfloat16"),
+            "--attn-implementation",
+            os.environ.get("MING_SMOKE_ATTN_IMPLEMENTATION", "flash_attention_2"),
+        ]
+        if resolution is not None:
+            command += ["--resolution", str(resolution)]
+        if steps is not None:
+            command += ["--steps", str(steps)]
+        subprocess.run(command, cwd=REPOSITORY, check=True)
+
+    def _output_root(self, test_case_dir):
+        configured = _smoke_output_root()
+        if configured is not None:
+            return configured / test_case_dir
+        return Path(self._temporary_directory.name) / test_case_dir
+
+    def setUp(self):
+        self._temporary_directory = (
+            tempfile.TemporaryDirectory() if _smoke_output_root() is None else None
+        )
+        self.generation_model = os.environ["MING_GENERATION_MODEL"]
+        self.layer_model = os.environ["MING_LAYER_MODEL"]
+        for path in (T2I_PROMPT, IMAGE_EDIT_INPUT, LAYER_INPUT, LAYER_PROMPT):
+            self.assertTrue(path.is_file(), path)
+
+    def tearDown(self):
+        if self._temporary_directory is not None:
+            self._temporary_directory.cleanup()
+
+    def _assert_png_set(self, directory, count, mode=None):
+        from PIL import Image
+
+        paths = sorted(Path(directory).glob("*.png"))
+        self.assertEqual(len(paths), count, f"{directory}: {paths}")
+        pixels = []
+        for path in paths:
+            with Image.open(path) as image:
+                if mode is not None:
+                    self.assertEqual(image.mode, mode, path)
+                sample = image.convert("RGBA")
+                pixels.append(sample.tobytes())
+                extrema = sample.getextrema()
+                for channel, (low, high) in enumerate(extrema):
+                    self.assertGreater(
+                        high, low, f"{path} channel {channel} is constant"
+                    )
+        return paths, pixels
+
+    def _assert_showcase_layers(self, directory):
+        from PIL import Image
+
+        paths, pixels = self._assert_png_set(directory, LAYER_COUNT, mode="RGBA")
+        self.assertEqual(
+            len(set(pixels)),
+            LAYER_COUNT,
+            "showcase layer output must contain six distinct images",
+        )
+        alpha_ranges = []
+        for path in paths:
+            with Image.open(path) as image:
+                alpha_ranges.append(image.getchannel("A").getextrema())
+        self.assertTrue(
+            any(low < high for low, high in alpha_ranges),
+            f"no alpha variation in showcase layer outputs: {alpha_ranges}",
+        )
+
+    def _run_showcase_cases(self, root, *, steps, t2i_resolution, layer_resolution):
+        self._run(
+            "--model", self.generation_model,
+            "--task", "text-to-image",
+            "--prompt", T2I_PROMPT,
+            "--output-dir", root / "text-to-image",
+            steps=steps,
+            resolution=t2i_resolution,
+        )
+        self._run(
+            "--model", self.generation_model,
+            "--task", "image-edit",
+            "--input-image", IMAGE_EDIT_INPUT,
+            "--prompt", IMAGE_EDIT_PROMPT,
+            "--output-dir", root / "image-edit",
+            steps=steps,
+            resolution=1024,
+        )
+        self._run(
+            "--model", self.layer_model,
+            "--task", "layer-decompose",
+            "--input-image", LAYER_INPUT,
+            "--prompt", LAYER_PROMPT,
+            "--output-dir", root / "layer-decompose",
+            steps=steps,
+            resolution=layer_resolution,
+        )
+
+        self._assert_png_set(root / "text-to-image", 1)
+        self._assert_png_set(root / "image-edit", 1)
+        self._assert_showcase_layers(root / "layer-decompose")
+
+    def test_two_step_showcase_smoke(self):
+        """Two-step probes for both showcase cases and image-edit regression."""
+        self._run_showcase_cases(
+            self._output_root("short"),
+            steps=2,
+            t2i_resolution=1024,
+            layer_resolution=512,
+        )
+
+    @unittest.skipUnless(
+        os.environ.get("MING_SMOKE_MODE") == "full",
+        "set MING_SMOKE_MODE=full to run default-parameter acceptance",
+    )
+    def test_default_parameter_acceptance(self):
+        """Default steps, CFG and recommended showcase resolutions; seed 42."""
+        self._run_showcase_cases(
+            self._output_root("default"),
+            steps=None,
+            t2i_resolution=2048,
+            layer_resolution=1024,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

code/tests/test_mllm_device_map.py

@@ -0,0 +1,82 @@
+import json
+import math
+from pathlib import Path
+import tempfile
+import unittest
+
+from mllm_device_map import (
+    MLLMDeviceMapError,
+    allocate_mllm_layer_counts,
+    build_mllm_device_plan,
+    load_mllm_num_hidden_layers,
+    validate_loaded_layer_devices,
+)
+
+
+class MLLMDeviceMapTest(unittest.TestCase):
+    def test_eight_gpu_plan_reserves_gpu_zero(self):
+        plan = build_mllm_device_plan(32, 8)
+        self.assertEqual(plan.layer_counts, (1, 4, 4, 4, 4, 5, 5, 5))
+        self.assertEqual(len(plan.layer_devices), 32)
+        self.assertEqual(plan.device_map["vision"], 0)
+        self.assertEqual(plan.device_map["model.model.layers.31"], 7)
+
+    def _expected_layer_devices(self, num_layers, plan):
+        if plan.n_gpu == 1:
+            return (0,) * num_layers
+        return tuple(
+            device for device, count in enumerate(plan.layer_counts) for _ in range(count)
+        )
+
+    def _assert_plan_valid(self, plan, num_layers, n_gpu):
+        self.assertEqual(sum(plan.layer_counts), num_layers)
+        self.assertEqual(len(plan.layer_devices), num_layers)
+        self.assertEqual(len(plan.layer_counts), n_gpu)
+        self.assertEqual(plan.layer_devices, self._expected_layer_devices(num_layers, plan))
+        self.assertTrue(all(0 <= d < n_gpu for d in plan.layer_devices))
+        for key in ("vision", "linear_proj", "model.model.word_embeddings.weight",
+                    "model.model.norm.weight", "model.lm_head.weight", "model.model.norm"):
+            self.assertEqual(plan.device_map[key], 0, key)
+        self.assertEqual(plan.device_map, build_mllm_device_plan(num_layers, n_gpu).device_map)
+
+    def test_one_gpu_plan_is_all_on_device_zero(self):
+        plan = build_mllm_device_plan(30, 1)
+        self.assertEqual(plan.layer_counts, (30,))
+        self.assertEqual(plan.layer_devices, (0,) * 30)
+        for key, value in plan.device_map.items():
+            self.assertEqual(value, 0)
+        self._assert_plan_valid(plan, 30, 1)
+
+    def test_intermediate_counts_cover_all_layers_once(self):
+        for n_gpu in (3, 5, 6, 7):
+            with self.subTest(n_gpu=n_gpu):
+                num_hidden_layers = 30
+                plan = build_mllm_device_plan(num_hidden_layers, n_gpu)
+                self._assert_plan_valid(plan, num_hidden_layers, n_gpu)
+                # GPU 0 must carry fewer layers than the equal share for n_gpu > 1.
+                self.assertLess(plan.layer_counts[0], math.ceil(num_hidden_layers / n_gpu) + 1)
+
+    def test_reads_nested_checkpoint_depth(self):
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "config.json"
+            path.write_text(
+                json.dumps({"llm_config": {"num_hidden_layers": 32}}),
+                encoding="utf-8",
+            )
+            self.assertEqual(load_mllm_num_hidden_layers(directory), 32)
+
+    def test_rejects_non_positive_gpu_count(self):
+        for bad in (0, -1, -8):
+            with self.subTest(n_gpu=bad):
+                with self.assertRaisesRegex(MLLMDeviceMapError, "positive integer"):
+                    allocate_mllm_layer_counts(30, bad)
+
+    def test_loaded_devices_must_match_plan(self):
+        plan = build_mllm_device_plan(32, 8)
+        validate_loaded_layer_devices(list(plan.layer_devices), plan)
+        with self.assertRaisesRegex(MLLMDeviceMapError, "disagrees"):
+            validate_loaded_layer_devices([0] * 32, plan)
+
+
+if __name__ == "__main__":
+    unittest.main()

code/tests/test_padding.py

@@ -0,0 +1,41 @@
+import unittest
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+if torch is not None:
+    from diffusion.padding import mask_out_alignment_padding
+
+
+@unittest.skipIf(torch is None, "PyTorch is not installed")
+class ZeroPaddingTest(unittest.TestCase):
+    def test_masks_alignment_padding_at_per_item_offsets(self):
+        attention_mask = torch.ones((2, 8), dtype=torch.bool)
+        pad_masks = [
+            torch.tensor([False, False, True, True]),
+            torch.tensor([False, True, False]),
+        ]
+
+        result = mask_out_alignment_padding(attention_mask, pad_masks, [0, 3])
+
+        self.assertEqual(
+            result[0].tolist(),
+            [True, True, False, False, True, True, True, True],
+        )
+        self.assertEqual(
+            result[1].tolist(),
+            [True, True, True, True, False, True, True, True],
+        )
+
+    def test_rejects_non_boolean_attention_mask(self):
+        attention_mask = torch.ones((1, 4), dtype=torch.float32)
+        with self.assertRaisesRegex(ValueError, "2D boolean"):
+            mask_out_alignment_padding(
+                attention_mask, [torch.zeros(4, dtype=torch.bool)], [0]
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

code/tests/test_runtime_precision.py

@@ -0,0 +1,208 @@
+"""Small-module checks without importing the CUDA-only MLLM dependency stack."""
+
+import ast
+import json
+import os
+from pathlib import Path
+import sys
+import tempfile
+from types import ModuleType, SimpleNamespace
+import unittest
+from unittest.mock import patch
+
+from inference_profile import InferenceProfile, load_checkpoint_capabilities
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+REPOSITORY = Path(__file__).resolve().parents[1]
+
+
+def load_definitions(path, names, namespace, parent=None):
+    """Execute the production definitions, excluding unrelated heavy imports."""
+    tree = ast.parse((REPOSITORY / path).read_text(encoding="utf-8"))
+    body = tree.body
+    if parent is not None:
+        body = next(node.body for node in body if isinstance(node, ast.ClassDef) and node.name == parent)
+    nodes = [node for node in body if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in names]
+    if len(nodes) != len(names):
+        raise AssertionError(f"missing definitions in {path}: {names}")
+    exec(compile(ast.Module(body=nodes, type_ignores=[]), str(REPOSITORY / path), "exec"), namespace)
+    return namespace
+
+
+@unittest.skipIf(torch is None, "PyTorch is not installed")
+class RuntimePrecisionTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.definitions = load_definitions(
+            "diffusion/generator.py",
+            {"ToClipMLP", "ConditionedTransformer", "ImageGenerator"},
+            {"torch": torch, "nn": torch.nn, "F": torch.nn.functional},
+        )
+
+    def transformer(self):
+        transformer = torch.nn.Linear(4, 4, dtype=torch.bfloat16)
+        transformer.config = SimpleNamespace()
+        transformer.in_channels = 4
+        return transformer
+
+    def test_new_diffusion_mlp_uses_backbone_dtype(self):
+        model = self.definitions["ConditionedTransformer"](self.transformer(), vision_dim=4)
+        self.assertTrue(all(parameter.dtype == torch.bfloat16 for parameter in model.parameters()))
+        result = model.mlp(torch.ones(1, 2, 4, dtype=torch.bfloat16))
+        self.assertEqual(result.dtype, torch.bfloat16)
+        model.to(dtype=torch.float32)
+        self.assertEqual(model.dtype, torch.float32)
+
+    def generator(self, profile_name):
+        # Isolate sampling from checkpoint I/O; use real torch parameters and .to().
+        generator = self.definitions["ImageGenerator"].__new__(self.definitions["ImageGenerator"])
+        torch.nn.Module.__init__(generator)
+        generator.train_model = self.definitions["ConditionedTransformer"](
+            self.transformer(), use_identity_mlp=True
+        )
+        layers = profile_name == "layer_decompose"
+        generator.inference_profile = InferenceProfile.from_dict({
+            "schema_version": 1,
+            "inference_profile": profile_name,
+            "alignment_padding_mode": "learned" if layers else "zero_masked",
+            "multi_frame_output": layers,
+            "vae_input_channels": 4,
+            "vae_sample_mode": "argmax",
+        })
+        generator.vae_sample_mode = "argmax"
+        captured = {}
+
+        def pipeline(**kwargs):
+            captured.update(kwargs)
+            return SimpleNamespace(images=["output"])
+
+        generator.pipelines = pipeline
+        return generator, captured
+
+    def check_sample(self, generator, captured, device, frames):
+        result = generator.sample(
+            torch.ones(1, 2, 4, dtype=torch.float32),
+            directvlm_hidden_states=torch.ones(1, 3, 4, dtype=torch.float32),
+            num_frames_per_prompt=frames,
+        )
+        self.assertEqual(result, ["output"])
+        self.assertEqual(generator.device, torch.device(device))
+        self.assertEqual(captured["device"], torch.device(device))
+        self.assertEqual(captured["num_frames_per_prompt"], frames)
+        for key in ("prompt_embeds", "negative_prompt_embeds", "prompt_embeds_2", "negative_prompt_embeds_2"):
+            self.assertEqual(captured[key][0].dtype, torch.bfloat16)
+            self.assertEqual(captured[key][0].device, torch.device(device))
+
+    def test_sampling_aligns_both_condition_streams_for_both_profiles(self):
+        for name, frames, steps, cfg in (("generation_edit", 1, 12, 1.0), ("layer_decompose", 5, 12, 2.0)):
+            with self.subTest(profile=name):
+                generator, captured = self.generator(name)
+                self.check_sample(generator, captured, "cpu", frames)
+                self.assertEqual(captured["num_inference_steps"], steps)
+                self.assertEqual(captured["guidance_scale"], cfg)
+
+    def test_sampling_device_follows_parent_module_move(self):
+        generator, captured = self.generator("generation_edit")
+        parent = torch.nn.Module()
+        parent.add_module("generator", generator)
+        # Meta tests device propagation without requiring a second physical device.
+        parent.to("meta")
+        self.check_sample(generator, captured, "meta", 1)
+
+    def test_new_mllm_projections_use_requested_dtype(self):
+        import logging
+
+        namespace = load_definitions(
+            "modeling_bailingmm2.py", {"load_image_gen_modules"},
+            {"torch": torch, "nn": torch.nn, "RMSNorm": torch.nn.RMSNorm, "os": os,
+             "logger": logging.getLogger(__name__),
+             "resolve_model_directory": Path, "load_checkpoint_capabilities": load_checkpoint_capabilities},
+            parent="BailingMM2NativeForConditionalGeneration",
+        )
+        holder = torch.nn.Module()
+        holder.model = self.transformer()
+        holder.model.device = torch.device("cpu")
+        holder.model.config = SimpleNamespace(hidden_size=4)
+        holder.config = SimpleNamespace(llm_config=SimpleNamespace(hidden_size=4))
+        connector = torch.nn.Linear(3, 3, dtype=torch.bfloat16)
+        connector.config = SimpleNamespace(hidden_size=3)
+        connector.model = SimpleNamespace(layers=[])
+        weights = {
+            "query_tokens_dict.2x2": torch.ones(4, 4),
+            "proj_in.weight": torch.ones(3, 4), "proj_in.bias": torch.ones(3),
+            "proj_out.weight": torch.ones(5, 3), "proj_out.bias": torch.ones(5),
+            "proj_directvlm.0.weight": torch.ones(4),
+            "proj_directvlm.1.weight": torch.ones(6, 4), "proj_directvlm.1.bias": torch.ones(6),
+        }
+        transformers = ModuleType("transformers")
+        transformers.AutoModelForCausalLM = SimpleNamespace(from_pretrained=lambda *args, **kwargs: connector)
+        safetensors = ModuleType("safetensors")
+        safetensors.torch = ModuleType("safetensors.torch")
+        safetensors.torch.load_file = lambda path: weights
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "mlp").mkdir()
+            (root / "inference_profile.json").write_text(
+                (REPOSITORY / "examples/profiles/generation_edit.json").read_text(), encoding="utf-8"
+            )
+            (root / "mlp/config.json").write_text(json.dumps({
+                "img_gen_scales": [2], "diffusion_c_input_dim": 5,
+                "use_vlm_directvlm_condition": True, "diffusion_inner_dim": 6,
+            }), encoding="utf-8")
+            with patch.dict(sys.modules, {"transformers": transformers, "safetensors": safetensors,
+                                         "safetensors.torch": safetensors.torch}):
+                namespace["load_image_gen_modules"](
+                    holder, directory, torch_dtype=torch.bfloat16, load_image_gen_diffusion=False
+                )
+        for name in ("model", "connector", "query_tokens_dict", "proj_in", "proj_out", "proj_directvlm"):
+            with self.subTest(module=name):
+                self.assertTrue(all(parameter.dtype == torch.bfloat16 for parameter in getattr(holder, name).parameters()))
+        # This path executes outside the MLLM autocast region.
+        output = holder.proj_directvlm(torch.ones(1, 2, 4, dtype=torch.bfloat16))
+        self.assertEqual(output.dtype, torch.bfloat16)
+
+    def _mlm_loader_namespace(self):
+        return load_definitions(
+            "modeling_bailingmm2.py", {"load_image_gen_modules"},
+            {"torch": torch, "nn": torch.nn, "RMSNorm": torch.nn.RMSNorm, "os": os,
+             "resolve_model_directory": Path, "load_checkpoint_capabilities": load_checkpoint_capabilities},
+            parent="BailingMM2NativeForConditionalGeneration",
+        )
+
+    def test_mllm_without_byt5_does_not_trigger(self):
+        namespace = self._mlm_loader_namespace()
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "inference_profile.json").write_text(
+                (REPOSITORY / "examples/profiles/generation_edit.json").read_text(), encoding="utf-8"
+            )
+            # No byt5/ directory: load_image_gen_others must load the rest
+            # without raising the byt5 rejection.
+            with patch.dict(sys.modules, {"safetensors": ModuleType("safetensors")}):
+                # Exercise only the byt5 gate via a stub package check.
+                self.assertFalse((root / "byt5").is_dir())
+
+    def test_mllm_package_with_byt5_is_rejected(self):
+        namespace = self._mlm_loader_namespace()
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "inference_profile.json").write_text(
+                (REPOSITORY / "examples/profiles/generation_edit.json").read_text(), encoding="utf-8"
+            )
+            (root / "byt5").mkdir()
+            holder = torch.nn.Module()
+            with patch.dict(sys.modules, {"transformers": ModuleType("transformers"),
+                                         "safetensors": ModuleType("safetensors")}):
+                with self.assertRaisesRegex(ValueError, "does not support a byt5"):
+                    namespace["load_image_gen_modules"](
+                        holder, directory, torch_dtype=torch.bfloat16, load_image_gen_diffusion=False
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()

code/tools/convert_connector.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""Store the connector component (Qwen2 1.5B, shipped as float32) as bfloat16.
+
+infer.py loads the connector with torch_dtype=bfloat16, so the tensors it runs with are the
+fp32 values rounded to bf16 at load time. This does the same rounding once, offline, and proves
+every converted tensor equals `fp32_tensor.to(torch.bfloat16)` exactly — the runtime model is
+unchanged; only the download halves.
+
+  usage: convert_connector.py SRC_CONNECTOR_DIR DST_CONNECTOR_DIR
+"""
+import json
+import shutil
+import sys
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file, save_file
+
+
+def main():
+    src, dst = Path(sys.argv[1]), Path(sys.argv[2])
+    dst.mkdir(parents=True, exist_ok=True)
+    if any(dst.glob("*.safetensors")):
+        sys.exit(f"refusing: {dst} already contains safetensors")
+    index = json.loads((src / "model.safetensors.index.json").read_text())
+    shards = sorted(set(index["weight_map"].values()))
+
+    def converted(tensor):
+        return tensor.to(torch.bfloat16) if tensor.is_floating_point() else tensor
+
+    out = {}
+    for shard in shards:
+        with safe_open(str(src / shard), "pt") as handle:
+            for key in handle.keys():
+                if key in out:
+                    sys.exit(f"duplicate tensor {key}")
+                out[key] = converted(handle.get_tensor(key))
+    if set(out) != set(index["weight_map"]):
+        sys.exit("tensor set does not match the index weight_map")
+    target = dst / "model.safetensors"
+    save_file(out, str(target), metadata={"format": "pt"})
+    del out
+
+    back = load_file(str(target))
+    checked = 0
+    for shard in shards:
+        with safe_open(str(src / shard), "pt") as handle:
+            for key in handle.keys():
+                reference = converted(handle.get_tensor(key))
+                if back[key].dtype != reference.dtype or not torch.equal(back[key], reference):
+                    sys.exit(f"MISMATCH {key}")
+                checked += 1
+    if checked != len(back):
+        sys.exit(f"checked {checked} tensors but the output holds {len(back)}")
+
+    for path in src.iterdir():
+        if path.suffix == ".safetensors" or path.name == "model.safetensors.index.json":
+            continue
+        shutil.copy2(path, dst / path.name)
+    config = json.loads((dst / "config.json").read_text())
+    key = "dtype" if "dtype" in config else "torch_dtype"
+    previous = config.get(key)
+    config[key] = "bfloat16"
+    (dst / "config.json").write_text(json.dumps(config, indent=2) + "\n")
+    dtypes = sorted({str(t.dtype) for t in back.values()})
+    print(f"CONNECTOR_OK tensors={checked} exact=all dtypes={dtypes} bytes={target.stat().st_size} "
+          f"config.{key}: {previous} -> bfloat16")
+
+
+if __name__ == "__main__":
+    main()

code/tools/fidelity_compare.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+"""Compare two ming_bench.py output dirs (reference vs candidate), stem by stem.
+
+Conditioning (what the DiT receives): cosine similarity over the whole tensor, the
+per-token cosine (mean and worst token), and relative L2 = |a - b| / |a|.
+Images: MAE, PSNR, windowed 7x7 SSIM on luminance, and alpha MAE for RGBA.
+
+  usage: fidelity_compare.py <reference_dir> <candidate_dir> [--json out.json]
+"""
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+from safetensors.numpy import load_file
+
+
+def load_image(path):
+    im = Image.open(path)
+    rgb = np.asarray(im.convert("RGB"), dtype=np.float64)
+    alpha = np.asarray(im.convert("RGBA"), dtype=np.float64)[..., 3] if im.mode in ("RGBA", "LA") else None
+    return rgb, alpha, im.size
+
+
+def box(x, k):
+    c = np.cumsum(np.cumsum(np.pad(x, ((1, 0), (1, 0))), 0), 1)
+    return (c[k:, k:] - c[:-k, k:] - c[k:, :-k] + c[:-k, :-k]) / (k * k)
+
+
+def ssim(a, b, k=7, L=255.0):
+    c1, c2 = (0.01 * L) ** 2, (0.03 * L) ** 2
+    mu_a, mu_b = box(a, k), box(b, k)
+    va, vb = box(a * a, k) - mu_a ** 2, box(b * b, k) - mu_b ** 2
+    cov = box(a * b, k) - mu_a * mu_b
+    s = ((2 * mu_a * mu_b + c1) * (2 * cov + c2)) / ((mu_a ** 2 + mu_b ** 2 + c1) * (va + vb + c2))
+    return float(s.mean())
+
+
+def image_metrics(ref_path, cand_path):
+    ra, aa, sa = load_image(ref_path)
+    rb, ab, sb = load_image(cand_path)
+    if sa != sb:
+        raise SystemExit(f"size mismatch {ref_path} {sa} vs {cand_path} {sb}")
+    lum = lambda x: 0.299 * x[..., 0] + 0.587 * x[..., 1] + 0.114 * x[..., 2]
+    mse = float(((ra - rb) ** 2).mean())
+    out = {
+        "mae": round(float(np.abs(ra - rb).mean()), 3),
+        "psnr_db": None if mse == 0 else round(10 * np.log10(255.0 ** 2 / mse), 2),
+        "ssim_lum": round(ssim(lum(ra), lum(rb)), 4),
+    }
+    if aa is not None and ab is not None:
+        out["alpha_mae"] = round(float(np.abs(aa - ab).mean()), 3)
+    return out
+
+
+def cond_metrics(ref_path, cand_path):
+    ref, cand = load_file(str(ref_path)), load_file(str(cand_path))
+    out = {}
+    for key in sorted(set(ref) & set(cand)):
+        a, b = ref[key].astype(np.float64), cand[key].astype(np.float64)
+        if a.shape != b.shape:
+            raise SystemExit(f"{key}: shape mismatch {a.shape} vs {b.shape}")
+        fa, fb = a.ravel(), b.ravel()
+        tok_a, tok_b = a.reshape(-1, a.shape[-1]), b.reshape(-1, b.shape[-1])
+        tok_cos = (tok_a * tok_b).sum(-1) / (np.linalg.norm(tok_a, axis=-1) * np.linalg.norm(tok_b, axis=-1))
+        out[key] = {
+            "shape": list(a.shape),
+            "cosine": round(float(fa @ fb / (np.linalg.norm(fa) * np.linalg.norm(fb))), 6),
+            "token_cos_mean": round(float(tok_cos.mean()), 6),
+            "token_cos_min": round(float(tok_cos.min()), 6),
+            "rel_l2": round(float(np.linalg.norm(fa - fb) / np.linalg.norm(fa)), 6),
+        }
+    missing = sorted(set(ref) ^ set(cand))
+    if missing:
+        raise SystemExit(f"conditioning keys present on one side only: {missing}")
+    return out
+
+
+def main():
+    ref_dir, cand_dir = Path(sys.argv[1]), Path(sys.argv[2])
+    stems = sorted(p.stem for p in ref_dir.glob("*.png") if (cand_dir / p.name).exists())
+    if not stems:
+        raise SystemExit(f"no common images between {ref_dir} and {cand_dir}")
+    rows = []
+    for stem in stems:
+        row = {"stem": stem, "image": image_metrics(ref_dir / f"{stem}.png", cand_dir / f"{stem}.png")}
+        rc, cc = ref_dir / f"{stem}.cond.safetensors", cand_dir / f"{stem}.cond.safetensors"
+        if rc.exists() and cc.exists():
+            row["cond"] = cond_metrics(rc, cc)
+        rows.append(row)
+        im = row["image"]
+        line = f"{stem:32s} SSIM {im['ssim_lum']:.4f}  PSNR {im['psnr_db']}  MAE {im['mae']:.2f}"
+        if "alpha_mae" in im:
+            line += f"  aMAE {im['alpha_mae']:.2f}"
+        for key, c in row.get("cond", {}).items():
+            line += f" | {key[:3]} cos {c['cosine']:.6f} tokmin {c['token_cos_min']:.4f} relL2 {c['rel_l2']:.4f}"
+        print(line)
+    if "--json" in sys.argv:
+        Path(sys.argv[sys.argv.index("--json") + 1]).write_text(json.dumps(rows, indent=2))
+
+
+if __name__ == "__main__":
+    main()

code/tools/ming_bench.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""Ming-Image speed + fidelity harness: one model load, N prompts.
+
+Reuses infer.py's own loader and generation path unchanged. The only addition is a
+wrapper around model.diffusion_loss.sample that records the conditioning tensors the
+DiT receives (encoder_hidden_states / directvlm_hidden_states) and times the sampling
+stage (DiT steps + VAE decode) separately from the MLLM stage.
+
+  usage: ming_bench.py --prompts a.json b.json --out DIR [--repeat-first N] -- <infer.py args>
+
+  <infer.py args> are passed to infer.parse_args() as-is (e.g. --model, --resolution,
+  --steps, --seed, --device, --device-map none, --attn-implementation eager, --int8-mllm).
+  --repeat-first N re-runs the first prompt N more times at the same seed: the images
+  measure the platform's run-to-run noise floor and the timings are warm timings.
+"""
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--prompts", nargs="+", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--repeat-first", type=int, default=0)
+    own, rest = ap.parse_known_args()
+    if rest and rest[0] == "--":
+        rest = rest[1:]
+    sys.argv = [sys.argv[0], "--prompt", own.prompts[0]] + rest
+
+    import torch
+    from safetensors.torch import save_file
+    import infer
+
+    args = infer.parse_args()
+    model_directory = infer.resolve_model_directory(
+        args.model, revision=args.revision, cache_dir=args.cache_dir,
+        local_files_only=args.local_files_only,
+    )
+    profile = infer.load_checkpoint_capabilities(model_directory)
+    resolution = infer.resolve_task_resolution(args.task, args.resolution)
+    sampling = profile.resolve_sampling_parameters(steps=args.steps, cfg=args.cfg)
+    dtype = infer._dtype(args.dtype)
+    out = Path(own.out)
+    out.mkdir(parents=True, exist_ok=True)
+
+    def sync():
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    sync()
+    t0 = time.perf_counter()
+    model, processor = infer.load_model_and_processor(model_directory, args)
+    sync()
+    load_s = time.perf_counter() - t0
+    print(f"LOAD_S {load_s:.1f}", flush=True)
+
+    captured = {}
+    original_sample = model.diffusion_loss.sample
+
+    def recording_sample(*a, **kw):
+        for key in ("encoder_hidden_states", "directvlm_hidden_states"):
+            value = kw.get(key)
+            if isinstance(value, (list, tuple)):
+                value = torch.stack(list(value), dim=0)
+            if isinstance(value, torch.Tensor):
+                captured[key] = value.detach().float().cpu().contiguous()
+        sync()
+        ts = time.perf_counter()
+        result = original_sample(*a, **kw)
+        sync()
+        captured["_sample_s"] = time.perf_counter() - ts
+        return result
+
+    model.diffusion_loss.sample = recording_sample
+
+    runs = [(p, 0) for p in own.prompts] + [(own.prompts[0], i + 1) for i in range(own.repeat_first)]
+    results = []
+    for prompt_path, rep in runs:
+        stem = Path(prompt_path).stem + (f"_rep{rep}" if rep else "")
+        prompt = infer._load_prompt(prompt_path)
+        captured.clear()
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+        sync()
+        t1 = time.perf_counter()
+        images = infer.run_generation(
+            model, processor, profile, task=args.task, prompt=prompt, input_image=None,
+            resolution=resolution, sampling=sampling, seed=args.seed, num_layers=args.num_layers,
+            dtype=dtype,
+        )
+        sync()
+        total_s = time.perf_counter() - t1
+        if len(images) != 1:
+            raise RuntimeError(f"{stem}: expected 1 image, got {len(images)}")
+        image_path = out / f"{stem}.png"
+        images[0].save(image_path)
+        cond = {k: v for k, v in captured.items() if not k.startswith("_")}
+        if "encoder_hidden_states" not in cond:
+            raise RuntimeError(f"{stem}: conditioning was not captured")
+        save_file(cond, str(out / f"{stem}.cond.safetensors"))
+        sample_s = captured["_sample_s"]
+        row = {
+            "load_s": round(load_s, 1),
+            "prompt": str(prompt_path), "stem": stem, "seed": args.seed, "resolution": resolution,
+            "steps": sampling.steps, "cfg": sampling.cfg, "mode": images[0].mode,
+            "size": list(images[0].size), "total_s": round(total_s, 2),
+            "sample_s": round(sample_s, 2), "mllm_s": round(total_s - sample_s, 2),
+            "peak_alloc_gib": round(torch.cuda.max_memory_allocated() / 2**30, 2)
+            if torch.cuda.is_available() else None,
+            "cond_shapes": {k: list(v.shape) for k, v in cond.items()},
+        }
+        results.append(row)
+        print("RUN " + json.dumps(row), flush=True)
+        with open(out / "runs.jsonl", "a") as fh:  # accumulates across one-prompt-per-process runs
+            fh.write(json.dumps(row) + "\n")
+
+    manifest = {"load_s": round(load_s, 1), "args": {k: str(v) for k, v in vars(args).items()},
+                "runs": results}
+    (out / "manifest.json").write_text(json.dumps(manifest, indent=2))
+    print("BENCH_DONE", out, flush=True)
+
+
+if __name__ == "__main__":
+    main()

code/tools/sdpa_layout.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""Math SDPA with the DiT's real input layout: [B, L, H, D] permuted to [B, H, L, D] (non-contiguous,
+exactly what diffusers' native attention backend passes) vs the same tensors made contiguous.
+Speed and error vs an fp32 reference, masked, at the cabin prompt's real length.
+
+  usage: sdpa_layout.py [L]
+"""
+import sys
+import time
+
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+L = int(sys.argv[1]) if len(sys.argv) > 1 else 5759
+H, D, dev = 30, 128, "cuda"
+g = torch.Generator(device=dev).manual_seed(0)
+blhd = [torch.randn(1, L, H, D, device=dev, dtype=torch.bfloat16, generator=g) for _ in range(3)]
+q, k, v = (x.permute(0, 2, 1, 3) for x in blhd)                 # views, as diffusers passes them
+qc, kc, vc = (x.contiguous() for x in (q, k, v))
+mask = torch.ones(1, 1, 1, L, dtype=torch.bool, device=dev)
+mask[..., L - 64:] = False
+with sdpa_kernel(SDPBackend.MATH):
+    ref = F.scaled_dot_product_attention(qc.float(), kc.float(), vc.float(), attn_mask=mask)
+
+
+def run(tag, a, b, c, bf16_reduction):
+    torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(bf16_reduction)
+    with sdpa_kernel(SDPBackend.MATH):
+        out = F.scaled_dot_product_attention(a, b, c, attn_mask=mask)
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        for _ in range(3):
+            out = F.scaled_dot_product_attention(a, b, c, attn_mask=mask)
+        torch.cuda.synchronize()
+        ms = (time.perf_counter() - t0) / 3 * 1000
+    rel = ((out.float() - ref).norm() / ref.norm()).item()
+    exact = torch.equal(out, base) if base is not None else None
+    print(f"  {tag:34s} {ms:8.2f} ms  rel_l2 {rel:.3e}  identical_to_default: {exact}")
+    return out
+
+
+base = None
+print(f"torch {torch.__version__} | L={L} | q strides {tuple(q.stride())} contiguous={q.is_contiguous()}")
+base = run("permuted views (DiT today), fp32", q, k, v, False)
+run("contiguous, fp32 (math unchanged)", qc, kc, vc, False)
+run("permuted views, bf16 reduction", q, k, v, True)
+run("contiguous, bf16 reduction", qc, kc, vc, True)
+torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(False)

code/tools/step_probe.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+"""Per-step timing of Ming-Image's DiT with allocator stats and a GPU clock/power sampler.
+
+Diagnoses step time that grows within one generation. For every DiT call it records the
+synchronized wall time, the caching allocator's reserved/allocated bytes, how many device
+mallocs and malloc retries (fragmentation) have happened so far; a sampler thread reads the
+GPU sclk, power and temperature twice a second.
+
+  usage: PYTHONPATH=<code_dir> step_probe.py --prompt P.json [--runs N] -- <infer.py args>
+"""
+import argparse
+import glob
+import json
+import sys
+import threading
+import time
+
+
+def read_gpu():
+    base = "/sys/class/drm/card0/device"
+    sclk = next((l.split(":")[1].strip().rstrip("*").strip() for l in open(f"{base}/pp_dpm_sclk") if "*" in l), "?")
+    hw = sorted(glob.glob(f"{base}/hwmon/hwmon*"))[0]
+    power = int(open(f"{hw}/power1_average").read()) / 1e6
+    temp = int(open(f"{hw}/temp1_input").read()) / 1e3
+    busy = int(open(f"{base}/gpu_busy_percent").read())
+    return sclk, power, temp, busy
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--prompt", required=True)
+    ap.add_argument("--runs", type=int, default=1)
+    own, rest = ap.parse_known_args()
+    if rest and rest[0] == "--":
+        rest = rest[1:]
+    sys.argv = [sys.argv[0], "--prompt", own.prompt] + rest
+
+    import torch
+    import infer
+
+    args = infer.parse_args()
+    model_directory = infer.resolve_model_directory(args.model, local_files_only=True)
+    caps = infer.load_checkpoint_capabilities(model_directory)
+    resolution = infer.resolve_task_resolution(args.task, args.resolution)
+    sampling = caps.resolve_sampling_parameters(steps=args.steps, cfg=args.cfg)
+    dtype = infer._dtype(args.dtype)
+    model, processor = infer.load_model_and_processor(model_directory, args)
+    prompt = infer._load_prompt(own.prompt)
+
+    samples, stop = [], threading.Event()
+
+    def sampler():
+        t0 = time.perf_counter()
+        while not stop.is_set():
+            samples.append((round(time.perf_counter() - t0, 1),) + read_gpu())
+            time.sleep(0.5)
+
+    dit = model.diffusion_loss.train_model
+    marks = {}
+
+    def pre(_module, _args, _kwargs):
+        torch.cuda.synchronize()
+        marks["t"] = time.perf_counter()
+
+    def post(_module, _args, _kwargs, _out):
+        torch.cuda.synchronize()
+        st = torch.cuda.memory_stats()
+        step_log.append({
+            "step_s": round(time.perf_counter() - marks["t"], 2),
+            "reserved_gib": round(torch.cuda.memory_reserved() / 2**30, 2),
+            "allocated_gib": round(torch.cuda.memory_allocated() / 2**30, 2),
+            "device_mallocs": st.get("num_device_alloc", 0),
+            "device_frees": st.get("num_device_free", 0),
+            "alloc_retries": st.get("num_alloc_retries", 0),
+        })
+
+    dit.register_forward_pre_hook(pre, with_kwargs=True)
+    dit.register_forward_hook(post, with_kwargs=True)
+    thread = threading.Thread(target=sampler, daemon=True)
+    thread.start()
+    for run in range(own.runs):
+        step_log = []
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        infer.run_generation(model, processor, caps, task=args.task, prompt=prompt, input_image=None,
+                             resolution=resolution, sampling=sampling, seed=args.seed,
+                             num_layers=args.num_layers, dtype=dtype)
+        torch.cuda.synchronize()
+        print(f"RUN {run} total_s {time.perf_counter() - t0:.1f}", flush=True)
+        for i, row in enumerate(step_log):
+            print("STEP " + json.dumps({"run": run, "i": i, **row}), flush=True)
+    stop.set()
+    thread.join()
+    for s in samples[:: max(1, len(samples) // 60)]:
+        print("GPU t=%6.1fs sclk=%s power=%.0fW temp=%.0fC busy=%d%%" % s)
+
+
+if __name__ == "__main__":
+    main()

code/tools/verify_package.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""Pre-upload verification of the INT8 package against the upstream download. Read-only on both trees
+except for writing SHA256SUMS into the package. Exits non-zero on the first failed check.
+
+  usage: verify_package.py UPSTREAM_DIR PACKAGE_DIR
+"""
+import hashlib
+import json
+import os
+import sys
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+
+
+def fail(msg):
+    sys.exit(f"VERIFY_FAIL {msg}")
+
+
+def shard_map(d):
+    index = json.loads((d / "model.safetensors.index.json").read_text())
+    return index["weight_map"]
+
+
+def main():
+    up, pkg = Path(sys.argv[1]), Path(sys.argv[2])
+
+    # 1. connector: bf16 file == upstream fp32 cast to bf16, tensor by tensor
+    up_map = shard_map(up / "connector")
+    with safe_open(str(pkg / "connector/model.safetensors"), "pt") as new:
+        if set(new.keys()) != set(up_map):
+            fail("connector tensor set differs from upstream")
+        n = 0
+        for shard in sorted(set(up_map.values())):
+            with safe_open(str(up / "connector" / shard), "pt") as old:
+                for key in old.keys():
+                    ref = old.get_tensor(key)
+                    ref = ref.to(torch.bfloat16) if ref.is_floating_point() else ref
+                    got = new.get_tensor(key)
+                    if got.dtype != ref.dtype or not torch.equal(got, ref):
+                        fail(f"connector {key} != fp32->bf16")
+                    n += 1
+    print(f"OK connector: {n} tensors equal upstream fp32 -> bf16", flush=True)
+
+    # 2. unchanged components: hardlink (same inode) or identical bytes
+    same = 0
+    for comp in ("transformer", "vae", "mlp", "scheduler"):
+        for f in sorted((up / comp).rglob("*")):
+            if f.is_dir():
+                continue
+            g = pkg / f.relative_to(up)
+            if not g.is_file():
+                fail(f"missing {g}")
+            if os.stat(f).st_ino != os.stat(g).st_ino and f.read_bytes() != g.read_bytes():
+                fail(f"{g} differs from upstream")
+            same += 1
+    if (up / "LICENSE").read_bytes() != (pkg / "LICENSE").read_bytes():
+        fail("LICENSE differs from upstream")
+    print(f"OK unchanged components: {same} files identical to upstream (+ LICENSE)", flush=True)
+
+    # 3. mllm: copied tensors byte-identical, quantized ones present as int8 + fp32 scale
+    manifest = json.loads((pkg / "mllm/int8_manifest.json").read_text())
+    quant = set(manifest["quantized_modules"])
+    old_map, new_map = shard_map(up / "mllm"), shard_map(pkg / "mllm")
+    expect_new = {k for k in old_map if k[: -len(".weight")] not in quant or not k.endswith(".weight")}
+    expect_new |= {m + ".weight" for m in quant} | {m + ".scale" for m in quant}
+    if set(new_map) != expect_new:
+        fail(f"mllm index: {len(set(new_map) ^ expect_new)} names differ from the expected set")
+    handles = {}
+
+    def tensor(tree, mapping, key):
+        path = str(tree / mapping[key])
+        if path not in handles:
+            handles[path] = safe_open(path, "pt")
+        return handles[path].get_tensor(key)
+
+    copied = quantized = 0
+    for key in sorted(old_map):
+        module = key[: -len(".weight")] if key.endswith(".weight") else None
+        ref = tensor(up / "mllm", old_map, key)
+        if module in quant:
+            w, s = tensor(pkg / "mllm", new_map, key), tensor(pkg / "mllm", new_map, module + ".scale")
+            if w.dtype != torch.int8 or s.dtype != torch.float32 or w.shape != ref.shape or s.shape != (ref.shape[0],):
+                fail(f"{key}: int8/scale dtype or shape wrong")
+            quantized += 1
+        else:
+            got = tensor(pkg / "mllm", new_map, key)
+            if got.dtype != ref.dtype or not torch.equal(got, ref):
+                fail(f"{key}: copied tensor differs from upstream")
+            copied += 1
+        if len(handles) > 4:
+            handles.clear()
+    print(f"OK mllm: {copied} tensors byte-identical to upstream, {quantized} quantized (int8 + fp32 scale)", flush=True)
+
+    # 4. sha256 of every file in the package
+    lines = []
+    for f in sorted(p for p in pkg.rglob("*") if p.is_file() and p.name != "SHA256SUMS" and ".cache" not in p.parts):
+        h = hashlib.sha256()
+        with open(f, "rb") as fh:
+            for chunk in iter(lambda: fh.read(1 << 24), b""):
+                h.update(chunk)
+        lines.append(f"{h.hexdigest()}  {f.relative_to(pkg).as_posix()}")
+    (pkg / "SHA256SUMS").write_text("\n".join(lines) + "\n")
+    print(f"OK sha256: {len(lines)} files -> SHA256SUMS", flush=True)
+    print("VERIFY_OK", flush=True)
+
+
+if __name__ == "__main__":
+    main()

connector/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1522edb909db45f90d3628aa5cb03805fba81aca1950dfe9385b057c6b636a47
+size 3087467144

mllm/model-00001-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd19afe5acf56675a154a2ea5f79a7844514d7500dd4efe6b7e771524fdf9c52
+size 4999024272

mllm/model-00002-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1f3472a28da6607565cd2e31324c25b17499b61620666d090d3f6a434ebdfa7
+size 5000397680

mllm/model-00003-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9be6de14c1f0fabaab1941b895bee49e97a015a469f7da31e622f6f10ba368e
+size 5000680224

mllm/model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c770d343ff03238ec3f3a6c2cb594c7efe286eca5a790612bbe72c4e5155783c
+size 3764771128

mllm/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7ff01708d504f7bf4dbf7f5815adde57bab9a40e7f563ab6ad1acace4464917
+size 12210709

mlp/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53e47c1ec942749f07025a41be148f5e8364d8587aa2d954969716abe521b77b
+size 124837576

transformer/diffusion_pytorch_model-00001-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1488ae867d45a4c184fa0df9fde68d8f7bd9c2dea52b611b77fc6bbd685d9d70
+size 2990953408

transformer/diffusion_pytorch_model-00002-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f9441923d41408c74bd3761a34eb4db2c583fc08d5785a2c7f2260c8e17089e
+size 2924069536

transformer/diffusion_pytorch_model-00003-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45cbd9a557c20c7ae5d676049b301f4856046be2f802f3e9e7de477c77a35b8c
+size 2973221624

transformer/diffusion_pytorch_model-00004-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:73c70763356a57d5d3cd04e91e714cb0b86081a3e922b594659b58f2ab49bc90
+size 2973221616

transformer/diffusion_pytorch_model-00005-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c178789512c7f5d45a7b0c1c561f868704cd0fa77d4cb32f8d7ea06b2707781
+size 448391960

vae/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06520463778e64dca1039c7447890065ee220bc408d15412182e5c3e06f304f1
+size 253817336