Add support for Ming-Image-0.1-Design with training adapter

README.md

@@ -35,6 +35,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [circlestone-labs/Anima-Base-v1.0-Diffusers](https://huggingface.co/circlestone-labs/Anima-Base-v1.0-Diffusers) (Anima)
 - [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (Krea 2)
 - [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (Krea 2 Turbo)
+- [inclusionAI/Ming-Image-0.1-Design](https://huggingface.co/inclusionAI/Ming-Image-0.1-Design) (Ming-Image 0.1 Design)
 - [microsoft/Mage-Flow-Base](https://huggingface.co/microsoft/Mage-Flow-Base) (Mage-Flow)
 
 ### Instruction / Edit

extensions_built_in/diffusion_models/__init__.py

@@ -6,6 +6,7 @@ from .flux_kontext import FluxKontextModel
 from .wan22 import Wan225bModel, Wan2214bModel, Wan2214bI2VModel
 from .qwen_image import QwenImageModel, QwenImageEditModel, QwenImageEditPlusModel
 from .qwen_image_2 import QwenImage2Model
+from .ming_image import MingImageModel
 from .flux2 import Flux2Model, Flux2Klein4BModel, Flux2Klein9BModel
 from .z_image import ZImageModel
 from .ltx2 import LTX2Model, LTX23Model, LTX25Model
@@ -38,6 +39,7 @@ AI_TOOLKIT_MODELS = [
     QwenImageEditModel,
     QwenImageEditPlusModel,
     QwenImage2Model,
+    MingImageModel,
     Flux2Model,
     ZImageModel,
     LTX2Model,

extensions_built_in/diffusion_models/ming_image/__init__.py

@@ -0,0 +1,3 @@
+from .ming_image import MingImageModel
+
+__all__ = ["MingImageModel"]

extensions_built_in/diffusion_models/ming_image/ming_image.py

@@ -0,0 +1,559 @@
+"""Ming-Image 0.1 Design for ai-toolkit.
+
+A 6B Z-Image-style DiT (`src/transformer.py`) driven by a Ling-mini-2.0 MoE
+multimodal LLM (`src/text_encoder.py`) and an RGBA Qwen-Image VAE. Text and
+an optional reference image go through the LLM; the DiT receives two caption
+streams (the 256 query-token states through the Qwen2 connector, and the
+"direct" LLM states of every prompt token) and, when editing, the reference's
+VAE latent as a second frame.
+
+Weights come from what `name_or_path` names (`src/checkpoints.py`): by
+default the ComfyUI single-file repack, whose int8_convrot files attach to the
+toolkit's convrot8 backend as-is (local copies win over the download), or the
+vendor's diffusers-style repo (`transformer/`, `vae/`, `mllm/`, `mlp/`,
+`connector/`) / a local checkpoint / a fine-tune loaded exactly as named. A
+full fine-tune saves the transformer as one .safetensors in the ComfyUI key
+layout, which ComfyUI and this class both load, the repack supplying the rest.
+
+Flow-matching convention matches ai-toolkit (t=1000 noise -> 0 clean, target =
+noise - clean); the DiT's own time/prediction convention is handled in
+`src/pipeline.py`.
+"""
+
+import os
+from typing import TYPE_CHECKING, List, Optional
+
+import huggingface_hub
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from safetensors.torch import load_file
+
+from toolkit.accelerator import unwrap_model
+from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
+from toolkit.basic import flush
+from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig
+from toolkit.lora_special import LoRASpecialNetwork
+from toolkit.metadata import get_meta_for_safetensors
+from toolkit.models.base_model import BaseModel
+from toolkit.samplers.custom_flowmatch_sampler import (
+    CustomFlowMatchEulerDiscreteScheduler,
+)
+
+from .src.checkpoints import BASE_REPO, COMFY_REPO
+from .src.pipeline import (
+    PIXELS_PER_TOKEN,
+    VAE_SCALE_FACTOR,
+    MingImagePipeline,
+    ming_shift_params,
+    run_transformer,
+)
+from .src.text_encoder import MingImagePromptEncoder, MingImageTextEncoder
+from .src.transformer import MingImageTransformer2DModel
+from .src.vae import MingImageVAE
+
+if TYPE_CHECKING:
+    from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
+
+
+# scheduler/scheduler_config.json ships shift=6.0 without dynamic shifting;
+# the reference pipeline forces dynamic shifting with the Z-Image defaults
+scheduler_config = {
+    "num_train_timesteps": 1000,
+    "shift": 6.0,
+    "use_dynamic_shifting": True,
+    "base_shift": 0.5,
+    "max_shift": 1.15,
+    "base_image_seq_len": 256,
+    "max_image_seq_len": 4096,
+    "time_shift_type": "exponential",
+}
+
+
+class MingImageFlowMatchScheduler(CustomFlowMatchEulerDiscreteScheduler):
+    """The toolkit scheduler with the reference pipeline's shift rule: the
+    max-shift endpoint moves with the image size (see `ming_shift_params`),
+    so 1024^2 and larger train at the same mu (1.35) they sample at."""
+
+    def set_train_timesteps(self, num_timesteps, device, timestep_type="linear", latents=None, patch_size=1):
+        if latents is not None and self.config.use_dynamic_shifting:
+            image_seq_len = latents.shape[2] * latents.shape[3] // (patch_size**2)
+            max_shift, max_seq_len = ming_shift_params(image_seq_len)
+            self.register_to_config(max_shift=max_shift, max_image_seq_len=max_seq_len)
+        return super().set_train_timesteps(
+            num_timesteps, device, timestep_type=timestep_type, latents=latents, patch_size=patch_size
+        )
+
+# decode above this many output pixels goes through the VAE's tiled path
+TILE_DECODE_ABOVE_PIXELS = 1024 * 1024
+
+
+class MingImageModel(BaseModel):
+    arch = "ming_image"
+
+    def __init__(
+        self,
+        device,
+        model_config: ModelConfig,
+        dtype="bf16",
+        custom_pipeline=None,
+        noise_scheduler=None,
+        **kwargs,
+    ):
+        super().__init__(
+            device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
+        )
+        self.is_flow_matching = True
+        self.is_transformer = True
+        self.use_old_lokr_format = False
+        self.target_lora_modules = ["MingImageTransformer2DModel"]
+        self.vae_scale_factor = VAE_SCALE_FACTOR
+        self.prompt_encoder: Optional[MingImagePromptEncoder] = None
+        self.image_processor = None
+
+        # Editing is what happens when the dataset has a control path: the one
+        # reference image rides into the LLM as vision tokens and its latent
+        # into the DiT as a second frame. The dataloader sizes it to the item's
+        # bucket, which is also the size the DiT needs it at.
+        self.encode_control_in_text_embeddings = True
+        self.has_multiple_control_images = False
+        self.use_raw_control_images = False
+
+    @staticmethod
+    def get_train_scheduler():
+        return MingImageFlowMatchScheduler(**scheduler_config)
+
+    def get_bucket_divisibility(self):
+        # 8 for the VAE, 2 for the DiT's 2x2 latent patches
+        return PIXELS_PER_TOKEN
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+    def load_model(self):
+        dtype = self.torch_dtype
+        self.print_and_status_update("Loading Ming-Image model")
+        model_path = self.model_config.name_or_path
+        extras_path = self.model_config.extras_name_or_path
+        if extras_path == model_path and model_path.endswith(".safetensors"):
+            # a single transformer file (a fine-tune): the rest comes from
+            # the repack
+            extras_path = COMFY_REPO
+        # only the vendor layout carries the configs, tokenizer and processor
+        config_path = extras_path if os.path.isdir(extras_path) else BASE_REPO
+
+        self.print_and_status_update("Loading transformer")
+        transformer = MingImageTransformer2DModel.load_model(
+            model_path,
+            dtype=dtype,
+            config_path=config_path,
+            use_comfy_weights=self.model_config.model_kwargs.get("use_comfy_weights", True),
+            qtype=self.model_config.qtype if self.model_config.quantize else None,
+            quantize_on_load=False,
+        )
+        load_kwargs = self.component_load_kwargs("transformer")
+        if self.model_config.assistant_lora_path is None:
+            transformer.aitk_post_load(**load_kwargs)
+        else:
+            # the adapter stays a live LoRA on the quantized linears, attached
+            # after quantization (it must wrap the quantized forward) and
+            # before layer offloading (the offloader routes through it)
+            offload = load_kwargs.pop("offload", 0.0)
+            transformer.aitk_post_load(offload=0.0, **load_kwargs)
+            self.load_training_adapter(transformer)
+            if offload and offload > 0:
+                from toolkit.memory_management import MemoryManager
+
+                MemoryManager.attach(
+                    transformer,
+                    torch.device(load_kwargs["quantize_device"]),
+                    offload_percent=offload,
+                    ignore_modules=list(transformer.get_offload_ignore_modules() or []),
+                )
+        flush()
+
+        self.print_and_status_update("Loading text encoder (MLLM + connector)")
+        tokenizer, image_processor = MingImageTextEncoder.load_tokenizer_and_processor(
+            config_path
+        )
+        text_encoder = MingImageTextEncoder.load(
+            extras_path, config_path=config_path, **self.component_load_kwargs("te")
+        )
+        text_encoder.requires_grad_(False)
+        text_encoder.eval()
+        flush()
+
+        self.print_and_status_update("Loading VAE")
+        vae = MingImageVAE.load(
+            extras_path, config_path=config_path, **self.component_load_kwargs("vae")
+        )
+        vae.requires_grad_(False)
+        vae.eval()
+
+        self.noise_scheduler = MingImageModel.get_train_scheduler()
+
+        self.vae = vae
+        self.text_encoder = [text_encoder]
+        self.tokenizer = [tokenizer]
+        self.image_processor = image_processor
+        self.model = transformer
+        self.prompt_encoder = MingImagePromptEncoder(text_encoder, tokenizer, image_processor)
+        self.pipeline = MingImagePipeline(self)
+        self.print_and_status_update("Model Loaded")
+
+    def load_training_adapter(self, transformer: MingImageTransformer2DModel):
+        """Attach the training adapter (a LoRA on the DiT) as a live assistant:
+        on at 1.0 while training, off while sampling so samples show the base
+        model. It is never merged: the weights are int8, whose grid swallows a
+        delta this small (~1e-4 against a ~6e-3 step) on a requantize."""
+        self.print_and_status_update("Loading assistant LoRA")
+        lora_path = self.model_config.assistant_lora_path
+        if not os.path.exists(lora_path):
+            # assume it is a hub path
+            lora_splits = lora_path.split("/")
+            if len(lora_splits) != 3:
+                raise ValueError(
+                    f"Assistant LoRA path {lora_path} is not a valid local path or hub path."
+                )
+            repo_id = "/".join(lora_splits[:2])
+            filename = lora_splits[2]
+            try:
+                lora_path = huggingface_hub.hf_hub_download(repo_id=repo_id, filename=filename)
+                self.model_config.assistant_lora_path = lora_path
+            except Exception as e:
+                raise ValueError(f"Failed to download assistant LoRA from {lora_path}: {e}")
+        lora_state_dict = load_file(lora_path)
+        dim = int(lora_state_dict["diffusion_model.layers.0.attention.to_k.lora_A.weight"].shape[0])
+        lora_state_dict = {
+            key.replace("diffusion_model.", "transformer."): value for key, value in lora_state_dict.items()
+        }
+
+        network_config = NetworkConfig(type="lora", linear=dim, linear_alpha=dim, transformer_only=True)
+        LoRASpecialNetwork.LORA_PREFIX_UNET = "lora_transformer"
+        network = LoRASpecialNetwork(
+            text_encoder=None,
+            unet=transformer,
+            lora_dim=network_config.linear,
+            multiplier=1.0,
+            alpha=network_config.linear_alpha,
+            train_unet=True,
+            train_text_encoder=False,
+            network_config=network_config,
+            network_type=network_config.type,
+            transformer_only=network_config.transformer_only,
+            is_transformer=True,
+            target_lin_modules=self.target_lora_modules,
+            is_assistant_adapter=True,
+            is_ara=True,
+            # transformer_only filters to the block list only through the holder
+            base_model=self,
+        )
+        network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True)
+        network.force_to(self.device_torch, dtype=self.torch_dtype)
+        network.eval()
+        network.requires_grad_(False)
+        network._update_torch_multiplier()
+        network.load_weights(lora_state_dict)
+        self.assistant_lora: LoRASpecialNetwork = network
+        self.assistant_lora.multiplier = 1.0
+        self.assistant_lora.is_active = True
+        self.invert_assistant_lora = False
+
+    # ------------------------------------------------------------------
+    # VAE. Four channels in and out: images without alpha get an opaque one on
+    # encode, and decode drops it again unless RGBA output is on.
+    # ------------------------------------------------------------------
+    @property
+    def load_rgba(self) -> bool:
+        return bool(self.model_config.model_kwargs.get("rgba", False))
+
+    def _to_rgba(self, images: torch.Tensor) -> torch.Tensor:
+        """`(B, 3|4, H, W)` in [-1, 1] -> four channels (opaque alpha added)."""
+        if images.shape[1] == 3:
+            images = torch.cat([images, torch.ones_like(images[:, :1])], dim=1)
+        return images
+
+    def encode_images(self, image_list, device=None, dtype=None):
+        device = device or self.vae_device_torch
+        dtype = dtype or self.vae_torch_dtype
+
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+        self.vae.eval()
+
+        images = torch.stack([image.to(device, dtype=dtype) for image in image_list])
+        images = self._to_rgba(images).unsqueeze(2)  # single-frame dim
+
+        # mode(), not sample(): this VAE's posterior std is ~1 in a latent
+        # space whose signal std is ~0.15, so sampling buries the image in noise
+        latents = self.vae.encode(images).latent_dist.mode()
+        latents = (latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+        return latents.squeeze(2).to(device, dtype=dtype)
+
+    def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
+        images = self._decode_rgba(latents, device=device, dtype=dtype)
+        if not self.load_rgba:
+            images = images[:, :3]
+        return images
+
+    def _decode_rgba(self, latents: torch.Tensor, device=None, dtype=None):
+        device = device or self.vae_device_torch
+        dtype = dtype or self.vae_torch_dtype
+
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+
+        latents = latents.to(device, dtype=dtype).unsqueeze(2)
+        latents = latents / self.vae.config.scaling_factor + self.vae.config.shift_factor
+
+        # the recommended 2048x2048 output is too much for a one-shot decode
+        # on a 32 GB card; tile above 1 MP, and whenever low_vram is set
+        pixels = latents.shape[-2] * latents.shape[-1] * self.vae_scale_factor**2
+        tiled = self.model_config.low_vram or pixels > TILE_DECODE_ABOVE_PIXELS
+        if tiled:
+            self.vae.enable_tiling(
+                tile_sample_min_height=1024,
+                tile_sample_min_width=1024,
+                tile_sample_stride_height=768,
+                tile_sample_stride_width=768,
+            )
+        try:
+            images = self.vae.decode(latents).sample
+        finally:
+            if tiled:
+                self.vae.disable_tiling()
+        return images.squeeze(2).to(device, dtype=dtype)
+
+    def decode_to_images(self, latents: torch.Tensor) -> List[Image.Image]:
+        """Decode to PIL, keeping the alpha channel when load_rgba is set."""
+        return [
+            self.image_tensor_to_pil(image) for image in self.decode_latents(latents)
+        ]
+
+    @staticmethod
+    def image_tensor_to_pil(image: torch.Tensor) -> Image.Image:
+        """`(C, H, W)` in [-1, 1] -> PIL (RGB, or RGBA when C is 4)."""
+        array = (image.float().clamp(-1, 1) / 2 + 0.5).permute(1, 2, 0)
+        array = (array.cpu().numpy() * 255).round().astype(np.uint8)
+        return Image.fromarray(array, mode="RGBA" if array.shape[2] == 4 else "RGB")
+
+    # ------------------------------------------------------------------
+    # Reference images
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _normalize_control_images(control_images, batch_size: int) -> List[List[torch.Tensor]]:
+        """Any shape the toolkit hands over -> one list (0 or 1 tensor, `(1, C, H, W)`
+        in [0, 1]) per batch item."""
+        if control_images is None:
+            return [[] for _ in range(batch_size)]
+        if isinstance(control_images, torch.Tensor):
+            if control_images.dim() == 5:
+                samples = [list(sample) for sample in control_images]
+            elif control_images.dim() == 4:
+                samples = [[sample] for sample in control_images]
+            else:
+                samples = [[control_images]]
+        elif len(control_images) > 0 and not isinstance(control_images[0], list):
+            samples = [list(control_images)]
+        else:
+            samples = [list(sample) for sample in control_images]
+        samples = [[img.unsqueeze(0) if img.dim() == 3 else img for img in sample] for sample in samples]
+        if len(samples) == 1 and batch_size > 1:
+            samples = samples * batch_size
+        if len(samples) != batch_size:
+            raise ValueError(f"got {len(samples)} control image sets for {batch_size} prompts")
+        for sample in samples:
+            if len(sample) > 1:
+                raise ValueError("Ming-Image takes a single reference image per prompt")
+        return samples
+
+    @staticmethod
+    def tensor_to_pil(image: torch.Tensor) -> Image.Image:
+        """`(1, C, H, W)` in [0, 1] -> RGB PIL for the vision tower (alpha dropped)."""
+        image = image[0, :3].detach().float().clamp(0, 1).cpu()
+        array = (image.permute(1, 2, 0).numpy() * 255).round().astype(np.uint8)
+        return Image.fromarray(array, mode="RGB")
+
+    def encode_reference_latents(self, images: List[Optional[torch.Tensor]], height: int, width: int):
+        """Per-item `(1, C, H, W)` references in [0, 1] (or None) -> per-item
+        `(C, h, w)` latents at the target size (or None)."""
+        out = []
+        for image in images:
+            if image is None:
+                out.append(None)
+                continue
+            image = image.to(self.device_torch, dtype=torch.float32)
+            if image.shape[-2:] != (height, width):
+                image = F.interpolate(image, size=(height, width), mode="bicubic", antialias=True).clamp(0, 1)
+            latent = self.encode_images(
+                [image[0] * 2 - 1], device=self.device_torch, dtype=self.torch_dtype
+            )
+            out.append(latent[0])
+        return out
+
+    # ------------------------------------------------------------------
+    # Prompts
+    # ------------------------------------------------------------------
+    def get_prompt_embeds(self, prompt, control_images=None) -> AdvancedPromptEmbeds:
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if self.text_encoder[0].device != self.device_torch:
+            self.text_encoder[0].to(self.device_torch)
+
+        samples = self._normalize_control_images(control_images, len(prompt))
+        images = [[self.tensor_to_pil(image) for image in sample] for sample in samples]
+
+        query, direct, has_reference = self.prompt_encoder.encode(prompt, images=images)
+        pe = AdvancedPromptEmbeds(
+            text_embeds=[q.to(self.torch_dtype) for q in query],
+            direct_embeds=[d.to(self.torch_dtype) for d in direct],
+            has_reference=[torch.tensor([flag], dtype=torch.bool) for flag in has_reference],
+        )
+        # bookkeeping, not activations: a .to(dtype) must not turn it into bf16
+        pe.frozen_dtype_keys = ["has_reference"]
+        return pe
+
+    # ------------------------------------------------------------------
+    # Training
+    # ------------------------------------------------------------------
+    def condition_noisy_latents(self, latents: torch.Tensor, batch: "DataLoaderBatchDTO"):
+        # the reference joins the DiT sequence in get_noise_prediction, clean
+        return latents.detach()
+
+    def get_noise_prediction(
+        self,
+        latent_model_input: torch.Tensor,  # (B, 16, h, w)
+        timestep: torch.Tensor,  # 0..1000 scale
+        text_embeddings: AdvancedPromptEmbeds,
+        batch: "DataLoaderBatchDTO" = None,
+        **kwargs,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+        batch_size = latent_model_input.shape[0]
+
+        query = list(text_embeddings.text_embeds)
+        direct = list(text_embeddings.direct_embeds)
+        flags = [bool(f.reshape(-1)[0]) for f in text_embeddings.has_reference]
+        if len(flags) == 1 and batch_size > 1:
+            flags = flags * batch_size
+
+        # The prompt decides: an embedding encoded with a reference expects the
+        # reference frame; one encoded without (plain T2I, a dropped caption)
+        # takes none, whatever the batch carries.
+        ref_latents = None
+        if batch is not None and any(flags):
+            control = batch.control_tensor_list
+            if control is None:
+                control = batch.control_tensor
+            if control is None:
+                raise ValueError(
+                    "the prompt embeddings were encoded with a reference image but the batch has none"
+                )
+            with torch.no_grad():
+                samples = self._normalize_control_images(control, batch_size)
+                images = [sample[0] if (flag and sample) else None for sample, flag in zip(samples, flags)]
+                ref_latents = self.encode_reference_latents(
+                    images,
+                    latent_model_input.shape[2] * VAE_SCALE_FACTOR,
+                    latent_model_input.shape[3] * VAE_SCALE_FACTOR,
+                )
+
+        return run_transformer(
+            self.transformer,
+            latent_model_input.to(self.device_torch, self.torch_dtype),
+            timestep,
+            query,
+            direct,
+            ref_latents=ref_latents,
+            **kwargs,
+        )
+
+    def get_loss_target(self, *args, **kwargs):
+        # flow-matching velocity target: noise - clean
+        noise = kwargs.get("noise")
+        batch = kwargs.get("batch")
+        return (noise - batch.latents).detach()
+
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    # ------------------------------------------------------------------
+    # Sampling
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return MingImagePipeline(self)
+
+    def generate_single_image(
+        self,
+        pipeline: MingImagePipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: AdvancedPromptEmbeds,
+        unconditional_embeds: AdvancedPromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        divisor = self.get_bucket_divisibility()
+        gen_config.width = int(gen_config.width // divisor * divisor)
+        gen_config.height = int(gen_config.height // divisor * divisor)
+
+        reference_latents = None
+        if gen_config.ctrl_img is not None:
+            # same channels the dataloader gives training references
+            mode = "RGBA" if self.load_rgba else "RGB"
+            image = (
+                torch.from_numpy(
+                    np.array(Image.open(gen_config.ctrl_img).convert(mode), dtype=np.float32) / 255.0
+                )
+                .permute(2, 0, 1)
+                .unsqueeze(0)
+            )
+            reference_latents = self.encode_reference_latents(
+                [image], gen_config.height, gen_config.width
+            )
+
+        return pipeline(
+            conditional_embeds=conditional_embeds,
+            unconditional_embeds=unconditional_embeds,
+            height=gen_config.height,
+            width=gen_config.width,
+            num_inference_steps=gen_config.num_inference_steps,
+            guidance_scale=gen_config.guidance_scale,
+            latents=gen_config.latents,
+            generator=generator,
+            reference_latents=reference_latents,
+        )[0]
+
+    # ------------------------------------------------------------------
+    # Saving / bookkeeping
+    # ------------------------------------------------------------------
+    def save_model(self, output_path, meta, save_dtype):
+        # one .safetensors in the ComfyUI key layout (fused qkv, root-level
+        # x_embedder / final_layer); pre-quantized layers keep their storage,
+        # so ComfyUI and this arch both load the file
+        transformer: MingImageTransformer2DModel = unwrap_model(self.model)
+        if not output_path.endswith(".safetensors"):
+            output_path += ".safetensors"
+        transformer.save_model(
+            output_path,
+            dtype=save_dtype,
+            metadata=get_meta_for_safetensors(meta, name=self.arch),
+        )
+
+    def get_base_model_version(self):
+        return self.arch
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return ["layers"]
+
+    def get_quantization_exclude_modules(self):
+        return MingImageTransformer2DModel.get_quantization_exclude_modules()
+
+    lora_keys_use_comfy_prefix = True

extensions_built_in/diffusion_models/ming_image/src/bailing_moe_v2.py

@@ -0,0 +1,613 @@
+"""BailingMoeV2 (Ling-mini-2.0) decoder for the Ming-Image conditioning stack.
+
+Ported from inclusionAI/Ming `modeling_bailing_moe_v2.py` (MIT) and cut down
+to the one thing Ming-Image asks of it: a single prefill pass over an already
+embedded sequence that returns the hidden states of selected layers. Kept:
+
+  - video-rope 3D positions (text tokens count along all three axes, image
+    and query-token blocks get t/h/w coordinates),
+  - partial rotary (half of each 128-wide head),
+  - the MultiRouter MoE: tokens that came in through the vision path route
+    through `image_gate`, everything else through `gate`,
+  - group-limited top-k routing with the expert bias and sigmoid scores.
+
+Dropped: KV cache and generation, the audio router, flash-attention paths.
+
+The routed experts of every MoE layer live in one stacked bank (`FusedExperts`)
+and run as grouped GEMMs (`moe_kernels.py`); the checkpoint's per-expert
+linears are stacked on load. The rope channel layout is built once per forward
+rather than once per layer.
+"""
+
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from toolkit.util.convrot_quant import largest_pow4_divisor, rotate
+
+from .moe_kernels import grouped_linear, grouped_swiglu
+
+
+class BailingMoeV2Config:
+    """The `llm_config` block of the checkpoint's mllm/config.json."""
+
+    def __init__(self, **kwargs):
+        self.vocab_size = 157184
+        self.hidden_size = 2048
+        self.intermediate_size = 5120
+        self.moe_intermediate_size = 512
+        self.num_hidden_layers = 20
+        self.num_attention_heads = 16
+        self.num_key_value_heads = 4
+        self.head_dim = 128
+        self.rms_norm_eps = 1e-6
+        self.rope_theta = 600000.0
+        self.partial_rotary_factor = 0.5
+        self.num_experts = 256
+        self.num_experts_per_tok = 8
+        self.num_shared_experts = 1
+        self.n_group = 8
+        self.topk_group = 4
+        self.routed_scaling_factor = 2.5
+        self.first_k_dense_replace = 1
+        self.use_qkv_bias = False
+        self.use_bias = False
+        self.router_type = "MultiRouter"
+        self.pad_token_id = 156892
+        self.image_patch_token = 157157
+        self.image_start_token = 157158
+        self.image_end_token = 157159
+        self.spatial_merge_size = 2
+        self.mrope_section = (8, 12, 12)
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+        if self.head_dim is None:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+
+
+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:
+        # normalize in fp32, apply the weight in the input dtype (as the original)
+        x = F.rms_norm(x.float(), (x.shape[-1],), eps=self.eps)
+        return self.weight * x.to(self.weight.dtype)
+
+
+class VideoRotaryEmbedding(nn.Module):
+    """cos/sin tables for 3D position ids, shaped `(3, B, L, rope_dim)`."""
+
+    def __init__(self, config: BailingMoeV2Config):
+        super().__init__()
+        dim = int(config.head_dim * config.partial_rotary_factor)
+        inv_freq = 1.0 / (
+            config.rope_theta
+            ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim)
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
+        inv_freq = (
+            self.inv_freq[None, None, :, None]
+            .float()
+            .expand(3, position_ids.shape[1], -1, 1)
+            .to(x.device)
+        )
+        positions = position_ids[:, :, None, :].float()
+        with torch.autocast(device_type=x.device.type, enabled=False):
+            freqs = (inv_freq.float() @ positions.float()).transpose(2, 3)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos, sin = emb.cos(), emb.sin()
+        return cos.to(x.dtype), sin.to(x.dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def video_rope_layout_index(rotary_dim: int, mrope_section: Sequence[int] = (8, 12, 12)) -> List[int]:
+    """Flat (axis * rotary_dim + channel) gather index realizing Ming's
+    `video_rope` channel layout: the h/w frequencies interleave channel by
+    channel and the t frequencies sit in their own runs, i.e.
+    (h,w,h,w,...)x24, t x8, (h,w,...)x24, t x8 for a 64-wide rotary half."""
+    section = [mrope_section[0], mrope_section[1] + mrope_section[2]] * 2
+    section = section[::-1]
+    rows = []
+    for i, width in enumerate(section):
+        if i % 2 == 0:
+            rows += [1 if j % 2 == 0 else 2 for j in range(width)]
+        else:
+            rows += [0] * width
+    # the layout is written for a 64-wide half; a narrower table (tiny test
+    # configs) keeps only the channels it has, as the slice loop it replaces did
+    return [row * rotary_dim + col for col, row in enumerate(rows) if col < rotary_dim]
+
+
+def video_rope_layout(cos, sin, flat: torch.Tensor):
+    """Apply the layout index once per forward. `cos`/`sin` are `(3, B, L,
+    rotary_dim)` tables per axis; returns `(B, 1, L, rotary_dim)` tables."""
+    # (3, B, L, D) -> (B, L, 3*D) -> pick (axis, channel) pairs -> (B, 1, L, D)
+    cos = cos.permute(1, 2, 0, 3).reshape(*cos.shape[1:3], -1).index_select(-1, flat)
+    sin = sin.permute(1, 2, 0, 3).reshape(*sin.shape[1:3], -1).index_select(-1, flat)
+    return cos.unsqueeze(1), sin.unsqueeze(1)
+
+
+def apply_video_rope(q, k, cos, sin):
+    """Rotate the first `cos.shape[-1]` channels of each head; `cos`/`sin`
+    come from `video_rope_layout`."""
+    rotary_dim = cos.shape[-1]
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+    q_rot = q_rot * cos + rotate_half(q_rot) * sin
+    k_rot = k_rot * cos + rotate_half(k_rot) * sin
+    return torch.cat([q_rot, q_pass], dim=-1), torch.cat([k_rot, k_pass], dim=-1)
+
+
+def get_video_rope_index(
+    config: BailingMoeV2Config,
+    input_ids: torch.Tensor,
+    attention_mask: torch.Tensor,
+    image_grid_thw: Optional[torch.Tensor],
+    scale_factor: float = 2.0,
+) -> torch.Tensor:
+    """3D position ids `(3, B, L)` for a text sequence with image blocks.
+
+    Text tokens advance all three axes together. Each image block (a
+    `<image>` start token followed by its patch tokens) gets t/h/w
+    coordinates centred on the block's row/column midpoints, offset by the
+    text position it follows; the text after it resumes from the block's
+    largest t. `image_grid_thw` lists the blocks in order of appearance
+    across the whole batch, in pre-merge patch units.
+    """
+    device = input_ids.device
+    if image_grid_thw is None or len(image_grid_thw) == 0:
+        position_ids = attention_mask.long().cumsum(-1) - 1
+        position_ids.masked_fill_(attention_mask == 0, 1)
+        return position_ids.unsqueeze(0).expand(3, -1, -1).to(device)
+
+    merge = config.spatial_merge_size
+    position_ids = torch.ones(
+        3, input_ids.shape[0], input_ids.shape[1], dtype=torch.long, device=device
+    )
+    image_index = 0
+    for b in range(input_ids.shape[0]):
+        valid = attention_mask[b] == 1
+        ids = input_ids[b][valid]
+        starts = torch.argwhere(ids == config.image_start_token).squeeze(1)
+        image_nums = int((ids[starts + 1] == config.image_patch_token).sum())
+        tokens = ids.tolist()
+        pos_list = []
+        st = 0
+        for _ in range(image_nums):
+            ed = tokens.index(config.image_patch_token, st)
+            t, h, w = (int(v) for v in image_grid_thw[image_index])
+            image_index += 1
+            grid_t, grid_h, grid_w = t, h // merge, w // merge
+            text_len = ed - st
+            st_idx = pos_list[-1][0].max() + 1 if pos_list else 0
+            pos_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+            t_index = torch.arange(grid_t).view(-1, 1).expand(-1, grid_h * grid_w).flatten()
+            h_index = (
+                torch.arange(grid_h).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten()
+                - (grid_h - 1) // 2
+            )
+            w_index = (
+                torch.arange(grid_w).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten()
+                - (grid_w - 1) // 2
+            )
+            t_index = t_index * scale_factor + text_len + st_idx
+            pos_list.append(torch.stack([t_index, h_index + t_index, w_index + t_index]))
+            st = ed + grid_t * grid_h * grid_w
+        if st < len(tokens):
+            st_idx = pos_list[-1][0].max() + 1 if pos_list else 0
+            pos_list.append(
+                torch.arange(len(tokens) - st).view(1, -1).expand(3, -1) + st_idx
+            )
+        llm_positions = torch.cat(pos_list, dim=1).reshape(3, -1).long()
+        position_ids[:, b, valid] = llm_positions.to(device)
+    return position_ids
+
+
+class MLP(nn.Module):
+    def __init__(self, config: BailingMoeV2Config, intermediate_size: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, 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 Gate(nn.Module):
+    """Sigmoid router with the load-balancing expert bias applied to the
+    selection only (weights come from the unbiased scores)."""
+
+    def __init__(self, config: BailingMoeV2Config):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.weight = nn.Parameter(torch.empty(config.num_experts, config.hidden_size))
+        self.expert_bias = nn.Parameter(torch.zeros(config.num_experts), requires_grad=False)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        num_tokens = x.shape[0]
+        logits = F.linear(x.float(), self.weight.float())
+        scores = torch.sigmoid(logits)
+        routing = scores + self.expert_bias.float()
+
+        group_scores = routing.view(num_tokens, self.n_group, -1).topk(2, dim=-1)[0].sum(-1)
+        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores).scatter_(1, group_idx, 1)
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(num_tokens, self.n_group, self.num_experts // self.n_group)
+            .reshape(num_tokens, -1)
+            .bool()
+        )
+        masked = routing.masked_fill(~score_mask, float("-inf"))
+        topk_idx = torch.topk(masked, k=self.top_k, dim=-1, sorted=False)[1]
+
+        topk_weight = torch.gather(scores, dim=1, index=topk_idx)
+        topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
+        return topk_idx, topk_weight * self.routed_scaling_factor
+
+
+class FusedExperts(nn.Module):
+    """All routed experts of one layer as stacked weight banks: `gate_up`
+    `(E, 2*inter, hidden)` (gate rows first) and `down` `(E, hidden, inter)`.
+    Runs as two grouped GEMMs over tokens sorted by expert.
+
+    Int8 storage is codes plus float32 per-output-channel scales (kept as uint8
+    byte views so `.to(dtype=...)` never casts them). `quantize_int8_` derives
+    it from the bf16 banks; `attach_int8_` takes it ready-made from a ComfyUI
+    checkpoint, whose convrot files hold the rows rotated by a block
+    regular-Hadamard along the input dim (`rot_size`). The rotation is
+    orthogonal and self-inverse, so the activations are rotated the same way
+    ahead of each GEMM instead of un-rotating the weights."""
+
+    def __init__(self, num_experts: int, hidden: int, inter: int):
+        super().__init__()
+        self.num_experts = num_experts
+        self.hidden = hidden
+        self.inter = inter
+        self.gate_up = nn.Parameter(torch.empty(num_experts, 2 * inter, hidden))
+        self.down = nn.Parameter(torch.empty(num_experts, hidden, inter))
+        self.is_int8 = False
+        self.rot_size = 1
+
+    def _rot(self, in_features: int) -> int:
+        # convrot's block: the group size, capped by the largest power of
+        # four dividing the input dim
+        return min(self.rot_size, largest_pow4_divisor(in_features))
+
+    @staticmethod
+    def fuse_state_dict(state: dict, prefix: str, num_experts: int) -> bool:
+        """Stack `{prefix}{e}.{gate,up,down}_proj.weight` entries into the bank
+        keys in place (the per-expert keys are removed). Returns whether any
+        per-expert keys were found."""
+        if f"{prefix}0.gate_proj.weight" not in state:
+            return False
+        gate_up = torch.stack(
+            [
+                torch.cat(
+                    [state.pop(f"{prefix}{e}.gate_proj.weight"), state.pop(f"{prefix}{e}.up_proj.weight")],
+                    dim=0,
+                )
+                for e in range(num_experts)
+            ]
+        )
+        down = torch.stack([state.pop(f"{prefix}{e}.down_proj.weight") for e in range(num_experts)])
+        state[f"{prefix}gate_up"] = gate_up
+        state[f"{prefix}down"] = down
+        return True
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        if f"{prefix}gate_up_q" in state_dict:
+            # ready-made int8 banks (a ComfyUI checkpoint, see
+            # MingImageTextEncoder.convert_state_dict_on_load)
+            rot = state_dict.pop(f"{prefix}rot_size", None)
+            self.attach_int8_(
+                state_dict.pop(f"{prefix}gate_up_q"),
+                state_dict.pop(f"{prefix}gate_up_scale"),
+                state_dict.pop(f"{prefix}down_q"),
+                state_dict.pop(f"{prefix}down_scale"),
+                rot_size=int(rot.item()) if rot is not None else 1,
+            )
+        else:
+            # accept the checkpoint's per-expert layout directly
+            self.fuse_state_dict(state_dict, prefix, self.num_experts)
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    @torch.no_grad()
+    def _set_int8(self, name: str, q: torch.Tensor, scale: torch.Tensor):
+        if name in self._parameters:
+            del self._parameters[name]
+        self.register_buffer(f"{name}_q", q.contiguous(), persistent=False)
+        self.register_buffer(
+            f"{name}_scale",
+            scale.detach().float().reshape(q.shape[0], q.shape[1]).contiguous().view(torch.uint8),
+            persistent=False,
+        )
+
+    @torch.no_grad()
+    def quantize_int8_(self, device=None):
+        """Per-output-channel symmetric int8 weights; the math runs on
+        `device` (the banks return to where they were)."""
+        if self.is_int8:
+            return
+        for name in ("gate_up", "down"):
+            param = self._parameters[name]
+            home = param.device
+            w = param.data.to(device or home, torch.float32)
+            scale = w.abs().amax(dim=-1).clamp_min(1e-12) / 127.0
+            q = torch.round(w / scale.unsqueeze(-1)).clamp_(-127, 127).to(torch.int8)
+            del w
+            self._set_int8(name, q.to(home), scale.to(home))
+        self.is_int8 = True
+        self.rot_size = 1
+
+    @torch.no_grad()
+    def attach_int8_(self, gate_up_q, gate_up_scale, down_q, down_scale, rot_size: int = 1):
+        """Ready-made int8 banks: codes `(E, N, K)` with float32 scales
+        `(E, N)` or `(E, N, 1)`; `rot_size` > 1 means the rows are rotated
+        along K in blocks of that size."""
+        expected = {
+            "gate_up": ((self.num_experts, 2 * self.inter, self.hidden), gate_up_q),
+            "down": ((self.num_experts, self.hidden, self.inter), down_q),
+        }
+        for name, (shape, q) in expected.items():
+            if tuple(q.shape) != shape or q.dtype != torch.int8:
+                raise ValueError(f"{name}: expected int8 bank {shape}, got {tuple(q.shape)} {q.dtype}")
+        self._set_int8("gate_up", gate_up_q, gate_up_scale)
+        self._set_int8("down", down_q, down_scale)
+        self.is_int8 = True
+        self.rot_size = int(rot_size)
+
+    @torch.no_grad()
+    def dequantize_(self, dtype=torch.bfloat16, device=None):
+        """Back to full-precision parameters (the rotation undone); the math
+        runs on `device`."""
+        if not self.is_int8:
+            return
+        for name in ("gate_up", "down"):
+            q = self._buffers.pop(f"{name}_q")
+            scale = self._buffers.pop(f"{name}_scale").view(torch.float32)
+            self._non_persistent_buffers_set.discard(f"{name}_q")
+            self._non_persistent_buffers_set.discard(f"{name}_scale")
+            home = q.device
+            work = device or home
+            w = q.to(work, torch.float32) * scale.to(work).unsqueeze(-1)
+            rot = self._rot(w.shape[-1])
+            if rot > 1:
+                w = rotate(w, rot)  # self-inverse
+            self.register_parameter(name, nn.Parameter(w.to(dtype).to(home), requires_grad=False))
+        self.is_int8 = False
+        self.rot_size = 1
+
+    def bank_names(self):
+        if self.is_int8:
+            return ("gate_up_q", "gate_up_scale", "down_q", "down_scale")
+        return ("gate_up", "down")
+
+    def bank_tensors(self) -> Dict[str, torch.Tensor]:
+        return {name: getattr(self, name) for name in self.bank_names()}
+
+    def _bank(self, tensors: Dict[str, torch.Tensor], name: str):
+        if self.is_int8:
+            return tensors[f"{name}_q"], tensors[f"{name}_scale"].view(torch.float32)
+        return tensors[name], None
+
+    def forward_with(self, x_sorted: torch.Tensor, offs: torch.Tensor, tensors: Dict[str, torch.Tensor]):
+        """The forward on an explicit set of bank tensors (the offload manager
+        hands over staged device copies)."""
+        if self.rot_size > 1:
+            x_sorted = rotate(x_sorted, self._rot(self.hidden))
+        h = grouped_swiglu(x_sorted, *self._bank(tensors, "gate_up"), offs)
+        if self.rot_size > 1:
+            h = rotate(h, self._rot(self.inter))
+        return grouped_linear(h, *self._bank(tensors, "down"), offs)
+
+    def forward(self, x_sorted: torch.Tensor, offs: torch.Tensor) -> torch.Tensor:
+        """`x_sorted` (rows, hidden) grouped by expert, `offs` (E + 1,) int32
+        prefix sums of the group sizes -> (rows, hidden)."""
+        return self.forward_with(x_sorted, offs, self.bank_tensors())
+
+
+class FusedExpertsMemoryManager:
+    """Layer offloading for an expert bank, in the shape of the toolkit's
+    per-layer memory managers: the banks stay pinned on cpu and are staged to
+    the compute device for each forward (0.75 GB int8 per layer, the bulk of
+    the model). `MemoryManager.detach` restores `_original_forward`."""
+
+    def __init__(self, module: FusedExperts, process_device: torch.device):
+        self.module = module
+        self.process_device = torch.device(process_device)
+        with torch.no_grad():
+            for name in module.bank_names():
+                tensor = getattr(module, name)
+                data = tensor.data if isinstance(tensor, nn.Parameter) else tensor
+                data = data.to("cpu")
+                if torch.cuda.is_available() and not data.is_pinned():
+                    try:
+                        data = data.pin_memory()
+                    except RuntimeError:
+                        pass
+                if isinstance(tensor, nn.Parameter):
+                    tensor.data = data
+                    tensor._is_memory_managed = True
+                else:
+                    module._buffers[name] = data
+        self._original_forward = module.forward
+
+        def _mm_forward(x_sorted, offs):
+            device = self.process_device
+            if device.type != "cuda":
+                return self._original_forward(x_sorted, offs)
+            staged = {
+                name: tensor.to(device, non_blocking=True)
+                for name, tensor in module.bank_tensors().items()
+            }
+            return module.forward_with(x_sorted.to(device), offs.to(device), staged)
+
+        module.forward = _mm_forward
+
+    @classmethod
+    def attach(cls, module: FusedExperts, process_device: torch.device):
+        if hasattr(module, "_layer_memory_manager"):
+            return
+        module._layer_memory_manager = cls(module, process_device)
+
+
+class SparseMoeBlock(nn.Module):
+    def __init__(self, config: BailingMoeV2Config):
+        super().__init__()
+        self.experts = FusedExperts(config.num_experts, config.hidden_size, config.moe_intermediate_size)
+        self.gate = Gate(config)
+        # MultiRouter: vision-path tokens have their own router
+        self.image_gate = Gate(config)
+        self.shared_experts = MLP(
+            config, config.moe_intermediate_size * config.num_shared_experts
+        )
+
+    def forward(self, hidden_states: torch.Tensor, image_mask: Optional[torch.Tensor]):
+        bsz, seq_len, dim = hidden_states.shape
+        x = hidden_states.reshape(-1, dim)
+        topk_idx, topk_weight = self.gate(x)
+        if image_mask is not None:
+            # both routers always run (the extra gate is one small matmul); no
+            # host sync on the mask
+            image_idx, image_weight = self.image_gate(x)
+            mask = image_mask.reshape(-1, 1)
+            topk_idx = torch.where(mask, image_idx, topk_idx)
+            topk_weight = torch.where(mask, image_weight, topk_weight)
+        y = self.moe_infer(x, topk_idx, topk_weight).view(bsz, seq_len, dim)
+        return y + self.shared_experts(hidden_states)
+
+    def moe_infer(self, x, topk_ids, topk_weight):
+        num_tokens, top_k = topk_ids.shape
+        flat_ids = topk_ids.reshape(-1)
+        order = flat_ids.argsort()
+        token_idx = order // top_k
+        # scatter_add rather than bincount: fixed output size, no host sync,
+        # graph-capturable
+        counts = torch.zeros(self.experts.num_experts, dtype=torch.int32, device=x.device)
+        counts.scatter_add_(0, flat_ids, torch.ones_like(flat_ids, dtype=torch.int32))
+        offs = F.pad(counts.cumsum(0), (1, 0)).to(torch.int32)
+        outs = self.experts(x[token_idx], offs)
+        weights = topk_weight.reshape(-1)[order].to(outs.dtype)
+        y = torch.zeros((num_tokens, x.shape[-1]), device=x.device, dtype=torch.float32)
+        y.index_add_(0, token_idx, (outs * weights.unsqueeze(-1)).float())
+        return y.to(x.dtype)
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return x
+    batch, heads, seq_len, head_dim = x.shape
+    return (
+        x[:, :, None, :, :]
+        .expand(batch, heads, n_rep, seq_len, head_dim)
+        .reshape(batch, heads * n_rep, seq_len, head_dim)
+    )
+
+
+class Attention(nn.Module):
+    def __init__(self, config: BailingMoeV2Config):
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.query_key_value = nn.Linear(
+            config.hidden_size,
+            (self.num_heads + 2 * self.num_kv_heads) * self.head_dim,
+            bias=config.use_qkv_bias,
+        )
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.dense = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=config.use_bias)
+
+    def forward(self, x, attn_mask, cos, sin):
+        bsz, seq_len, _ = x.shape
+        qkv = self.query_key_value(x).view(bsz, seq_len, self.num_heads + 2 * self.num_kv_heads, self.head_dim)
+        q, k, v = qkv.split([self.num_heads, self.num_kv_heads, self.num_kv_heads], dim=-2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+        q, k = self.q_norm(q), self.k_norm(k)
+        q, k = apply_video_rope(q, k, cos, sin)
+        k = repeat_kv(k, self.num_heads // self.num_kv_heads)
+        v = repeat_kv(v, self.num_heads // self.num_kv_heads)
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        return self.dense(out.transpose(1, 2).reshape(bsz, seq_len, -1))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: BailingMoeV2Config, layer_idx: int):
+        super().__init__()
+        self.attention = Attention(config)
+        self.is_moe = layer_idx >= config.first_k_dense_replace
+        self.mlp = SparseMoeBlock(config) if self.is_moe else MLP(config, config.intermediate_size)
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(self, x, attn_mask, cos, sin, image_mask):
+        x = x + self.attention(self.input_layernorm(x), attn_mask, cos, sin)
+        h = self.post_attention_layernorm(x)
+        h = self.mlp(h, image_mask) if self.is_moe else self.mlp(h)
+        return x + h
+
+
+class BailingMoeV2Model(nn.Module):
+    def __init__(self, config: BailingMoeV2Config):
+        super().__init__()
+        self.config = config
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        self.layers = nn.ModuleList(
+            [DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = VideoRotaryEmbedding(config)
+        # built once (a host->device copy would break CUDA-graph capture)
+        self.register_buffer(
+            "rope_layout_index",
+            torch.tensor(
+                video_rope_layout_index(self.rotary_emb.inv_freq.numel() * 2, config.mrope_section),
+                dtype=torch.long,
+            ),
+            persistent=False,
+        )
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,  # (B, L, hidden)
+        attention_mask: torch.Tensor,  # (B, 1, L, L) bool, True = may attend
+        position_ids: torch.Tensor,  # (3, B, L)
+        image_mask: Optional[torch.Tensor] = None,  # (B, L) bool, vision-path tokens
+        capture_layers: Sequence[int] = (),
+    ) -> Tuple[torch.Tensor, Dict[int, torch.Tensor]]:
+        """Returns the final (normalized) hidden states and, for each index in
+        `capture_layers`, the residual stream entering that layer; the index
+        equal to the layer count maps to the normalized output, matching the
+        `output_hidden_states` numbering of the reference implementation."""
+        cos, sin = self.rotary_emb(inputs_embeds, position_ids)
+        cos, sin = video_rope_layout(cos, sin, self.rope_layout_index)
+        capture = set(capture_layers)
+        captured: Dict[int, torch.Tensor] = {}
+        h = inputs_embeds
+        for i, layer in enumerate(self.layers):
+            if i in capture:
+                captured[i] = h
+            h = layer(h, attention_mask, cos, sin, image_mask)
+        h = self.norm(h)
+        if len(self.layers) in capture:
+            captured[len(self.layers)] = h
+        return h, captured

extensions_built_in/diffusion_models/ming_image/src/checkpoints.py

@@ -0,0 +1,34 @@
+"""Where Ming-Image weights come from.
+
+`name_or_path` is authoritative. COMFY_REPO, the single-file ComfyUI repack,
+is the default: its int8_convrot files are the toolkit's convrot8 storage, so a
+quantized load attaches them as-is, and copies already in the local ComfyUI
+models folder win over the download. Anything else in `name_or_path` (the
+vendor repo, a local checkpoint, a fine-tune) loads exactly what it names,
+local comfy files or not. BASE_REPO is the vendor's diffusers-layout
+checkpoint; its configs, tokenizer, image processor and (until the repack
+carries one) the vision tower are used with the repack. The repack tracks a
+work-in-progress ComfyUI PR and is expected to move: change COMFY_REPO and the
+UI default together.
+"""
+
+BASE_REPO = "inclusionAI/Ming-Image-0.1-Design"
+COMFY_REPO = "Kijai/Ming-Image-ComfyUI"
+
+# repo-relative comfy files per component; the resolver orders them by the
+# requested qtype (convrot8 -> the int8 file, anything else -> bf16)
+COMFY_TRANSFORMER_FILES = [
+    "diffusion_models/ming_image_0.1_design_int8_convrot.safetensors",
+    "diffusion_models/ming_image_0.1_design_bf16.safetensors",
+]
+COMFY_TEXT_ENCODER_FILES = [
+    "text_encoders/ming_image_0.1_ling_mini_2.0_int8_convrot.safetensors",
+    "text_encoders/ming_image_0.1_ling_mini_2.0_bf16.safetensors",
+]
+COMFY_VAE_FILES = ["vae/ming_image_vae_bf16.safetensors"]
+
+
+def comfy_weight_names(files):
+    """OstrisModelMixin candidate map: only the repack id resolves to these
+    files; every other name_or_path loads what it names."""
+    return {COMFY_REPO: list(files)}

extensions_built_in/diffusion_models/ming_image/src/moe_kernels.py

@@ -0,0 +1,264 @@
+"""Grouped GEMMs for the Ling-mini MoE experts.
+
+Tokens arrive sorted by expert with `offs` (num_experts + 1 int32 prefix sums)
+marking each expert's row range; every expert's weights sit in one stacked bank
+`(E, N, K)`. One launch per projection covers all experts, so a forward is a
+handful of kernels instead of one per (expert, linear). Weights are bf16 or
+int8 with a float32 per-output-channel scale (weight-only quantization; the
+activations stay bf16 and the accumulator is fp32).
+
+`grouped_swiglu(x, gate_up, scale, offs)` -> silu(x W_gate^T) * (x W_up^T)
+`grouped_linear(x, w, scale, offs)`       -> x W^T
+
+Both fall back to a per-expert torch loop without triton or off the gpu.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+try:
+    import triton
+    import triton.language as tl
+
+    HAS_TRITON = True
+except ImportError:  # pragma: no cover
+    HAS_TRITON = False
+
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _grouped_swiglu_kernel(
+        x_ptr,
+        w_ptr,
+        s_ptr,
+        offs_ptr,
+        y_ptr,
+        K: tl.constexpr,
+        N: tl.constexpr,
+        stride_xm,
+        stride_ym,
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        BLOCK_K: tl.constexpr,
+        W_INT8: tl.constexpr,
+    ):
+        e = tl.program_id(0)
+        nb = tl.program_id(1)
+        start = tl.load(offs_ptr + e)
+        end = tl.load(offs_ptr + e + 1)
+        if start == end:
+            return
+        offs_n = nb * BLOCK_N + tl.arange(0, BLOCK_N)
+        n_mask = offs_n < N
+        offs_k = tl.arange(0, BLOCK_K)
+        w_base = w_ptr + e.to(tl.int64) * (2 * N * K)
+        g_ptrs = w_base + offs_n[:, None] * K + offs_k[None, :]
+        u_ptrs = w_base + (offs_n[:, None] + N) * K + offs_k[None, :]
+        if W_INT8:
+            s_g = tl.load(s_ptr + e * (2 * N) + offs_n, mask=n_mask, other=0.0)
+            s_u = tl.load(s_ptr + e * (2 * N) + N + offs_n, mask=n_mask, other=0.0)
+        for m0 in range(start, end, BLOCK_M):
+            offs_m = m0 + tl.arange(0, BLOCK_M)
+            m_mask = offs_m < end
+            acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+            acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+            x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :]
+            for k0 in range(0, K, BLOCK_K):
+                a = tl.load(x_ptrs + k0, mask=m_mask[:, None], other=0.0)
+                g = tl.load(g_ptrs + k0, mask=n_mask[:, None], other=0)
+                u = tl.load(u_ptrs + k0, mask=n_mask[:, None], other=0)
+                if W_INT8:
+                    g = g.to(tl.bfloat16)
+                    u = u.to(tl.bfloat16)
+                acc_g = tl.dot(a, tl.trans(g), acc_g)
+                acc_u = tl.dot(a, tl.trans(u), acc_u)
+            if W_INT8:
+                acc_g = acc_g * s_g[None, :]
+                acc_u = acc_u * s_u[None, :]
+            h = acc_g * tl.sigmoid(acc_g) * acc_u
+            y_ptrs = y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :]
+            tl.store(y_ptrs, h.to(y_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :])
+
+    @triton.jit
+    def _grouped_linear_kernel(
+        x_ptr,
+        w_ptr,
+        s_ptr,
+        offs_ptr,
+        y_ptr,
+        K: tl.constexpr,
+        N: tl.constexpr,
+        stride_xm,
+        stride_ym,
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        BLOCK_K: tl.constexpr,
+        W_INT8: tl.constexpr,
+    ):
+        e = tl.program_id(0)
+        nb = tl.program_id(1)
+        start = tl.load(offs_ptr + e)
+        end = tl.load(offs_ptr + e + 1)
+        if start == end:
+            return
+        offs_n = nb * BLOCK_N + tl.arange(0, BLOCK_N)
+        n_mask = offs_n < N
+        offs_k = tl.arange(0, BLOCK_K)
+        w_ptrs = w_ptr + e.to(tl.int64) * (N * K) + offs_n[:, None] * K + offs_k[None, :]
+        if W_INT8:
+            s = tl.load(s_ptr + e * N + offs_n, mask=n_mask, other=0.0)
+        for m0 in range(start, end, BLOCK_M):
+            offs_m = m0 + tl.arange(0, BLOCK_M)
+            m_mask = offs_m < end
+            acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+            x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :]
+            for k0 in range(0, K, BLOCK_K):
+                a = tl.load(x_ptrs + k0, mask=m_mask[:, None], other=0.0)
+                w = tl.load(w_ptrs + k0, mask=n_mask[:, None], other=0)
+                if W_INT8:
+                    w = w.to(tl.bfloat16)
+                acc = tl.dot(a, tl.trans(w), acc)
+            if W_INT8:
+                acc = acc * s[None, :]
+            y_ptrs = y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :]
+            tl.store(y_ptrs, acc.to(y_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :])
+
+
+def _use_triton(x: torch.Tensor, K: int) -> bool:
+    # the kernels tile K by 64 without masking
+    return HAS_TRITON and x.is_cuda and K % 64 == 0
+
+
+def _block_m(rows: int, num_experts: int) -> int:
+    # average rows per expert decides the M tile; tiny groups waste less with 16
+    avg = rows / max(num_experts, 1)
+    if avg >= 96:
+        return 128
+    if avg >= 48:
+        return 64
+    if avg >= 24:
+        return 32
+    return 16
+
+
+def _check(x, w, scale, offs):
+    E, N, K = w.shape
+    if x.dim() != 2 or x.shape[1] != K:
+        raise ValueError(f"expected activations (rows, {K}), got {tuple(x.shape)}")
+    if offs.shape != (E + 1,):
+        raise ValueError(f"offs must have {E + 1} entries, got {tuple(offs.shape)}")
+    if w.dtype == torch.int8:
+        if scale is None or scale.shape != (E, N) or scale.dtype != torch.float32:
+            raise ValueError("int8 weights need a float32 (E, N) scale")
+    elif w.dtype != x.dtype:
+        raise ValueError(f"weight dtype {w.dtype} does not match activations {x.dtype}")
+
+
+def grouped_swiglu(
+    x: torch.Tensor, gate_up: torch.Tensor, scale: Optional[torch.Tensor], offs: torch.Tensor
+) -> torch.Tensor:
+    """`x` (rows, K) sorted by expert; `gate_up` (E, 2*I, K) with the gate rows
+    first; returns silu(gate) * up as (rows, I)."""
+    _check(x, gate_up, scale, offs)
+    E, two_i, K = gate_up.shape
+    inter = two_i // 2
+    if not _use_triton(x, K):
+        return _swiglu_fallback(x, gate_up, scale, offs)
+    x = x.contiguous()
+    y = torch.empty((x.shape[0], inter), device=x.device, dtype=x.dtype)
+    if x.shape[0] == 0:
+        return y
+    block_m = _block_m(x.shape[0], E)
+    block_n = 64
+    grid = (E, triton.cdiv(inter, block_n))
+    _grouped_swiglu_kernel[grid](
+        x,
+        gate_up,
+        scale if scale is not None else gate_up,
+        offs,
+        y,
+        K,
+        inter,
+        x.stride(0),
+        y.stride(0),
+        BLOCK_M=block_m,
+        BLOCK_N=block_n,
+        BLOCK_K=64,
+        W_INT8=gate_up.dtype == torch.int8,
+        num_warps=4,
+        num_stages=3,
+    )
+    return y
+
+
+def grouped_linear(
+    x: torch.Tensor, w: torch.Tensor, scale: Optional[torch.Tensor], offs: torch.Tensor
+) -> torch.Tensor:
+    """`x` (rows, K) sorted by expert; `w` (E, N, K); returns (rows, N)."""
+    _check(x, w, scale, offs)
+    E, N, K = w.shape
+    if not _use_triton(x, K):
+        return _linear_fallback(x, w, scale, offs)
+    x = x.contiguous()
+    y = torch.empty((x.shape[0], N), device=x.device, dtype=x.dtype)
+    if x.shape[0] == 0:
+        return y
+    block_m = _block_m(x.shape[0], E)
+    block_n = 64
+    grid = (E, triton.cdiv(N, block_n))
+    _grouped_linear_kernel[grid](
+        x,
+        w,
+        scale if scale is not None else w,
+        offs,
+        y,
+        K,
+        N,
+        x.stride(0),
+        y.stride(0),
+        BLOCK_M=block_m,
+        BLOCK_N=block_n,
+        BLOCK_K=64,
+        W_INT8=w.dtype == torch.int8,
+        num_warps=4,
+        num_stages=3,
+    )
+    return y
+
+
+# ---------------- torch fallbacks (cpu, no triton) ----------------
+
+
+def _expert_weight(w: torch.Tensor, scale: Optional[torch.Tensor], e: int, dtype: torch.dtype):
+    if w.dtype == torch.int8:
+        return (w[e].float() * scale[e][:, None]).to(dtype)
+    return w[e].to(dtype)
+
+
+def _swiglu_fallback(x, gate_up, scale, offs):
+    inter = gate_up.shape[1] // 2
+    y = torch.empty((x.shape[0], inter), device=x.device, dtype=x.dtype)
+    bounds = offs.tolist()
+    for e in range(gate_up.shape[0]):
+        start, end = bounds[e], bounds[e + 1]
+        if start == end:
+            continue
+        w = _expert_weight(gate_up, scale, e, torch.float32)
+        h = x[start:end].float() @ w.t()
+        y[start:end] = (F.silu(h[:, :inter]) * h[:, inter:]).to(x.dtype)
+    return y
+
+
+def _linear_fallback(x, w, scale, offs):
+    y = torch.empty((x.shape[0], w.shape[1]), device=x.device, dtype=x.dtype)
+    bounds = offs.tolist()
+    for e in range(w.shape[0]):
+        start, end = bounds[e], bounds[e + 1]
+        if start == end:
+            continue
+        we = _expert_weight(w, scale, e, torch.float32)
+        y[start:end] = (x[start:end].float() @ we.t()).to(x.dtype)
+    return y

extensions_built_in/diffusion_models/ming_image/src/pipeline.py

@@ -0,0 +1,164 @@
+"""Transformer call convention and a preview sampler for Ming-Image.
+
+Timesteps: the toolkit counts 0..1000 with 1000 = pure noise; the Z-Image DiT
+takes t in [0, 1] with 1 = clean and predicts `clean - noise`, so the call is
+`transformer(x, 1 - t/1000, ...)` and the toolkit's velocity is the negated
+output. Negative conditioning is all zeros (both streams), as in the reference
+implementation; guidance 1.0 (no CFG) is the recommended setting.
+"""
+
+from typing import List, Optional
+
+import numpy as np
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+from toolkit.samplers.custom_flowmatch_sampler import calculate_shift
+
+# Z-Image's dynamic shift (0.5 at 256 tokens -> 1.15 at 4096) below 1024^2;
+# the reference pipeline pins mu at 1.35 for 1024^2 and above instead of
+# extrapolating the line
+SHIFT_BASE_SEQ_LEN = 256
+SHIFT_MAX_SEQ_LEN = 4096
+SHIFT_BASE = 0.5
+SHIFT_MAX = 1.15
+SHIFT_MAX_LARGE = 1.35
+
+
+def ming_shift_params(image_seq_len: int):
+    """(max_shift, max_image_seq_len) the reference uses for this many image tokens."""
+    if image_seq_len >= SHIFT_MAX_SEQ_LEN:
+        return SHIFT_MAX_LARGE, image_seq_len
+    return SHIFT_MAX, SHIFT_MAX_SEQ_LEN
+
+
+def ming_shift_mu(image_seq_len: int) -> float:
+    max_shift, max_seq_len = ming_shift_params(image_seq_len)
+    return calculate_shift(image_seq_len, SHIFT_BASE_SEQ_LEN, max_seq_len, SHIFT_BASE, max_shift)
+
+VAE_SCALE_FACTOR = 8
+# 8 for the VAE, 2 for the DiT's 2x2 latent patches
+PIXELS_PER_TOKEN = VAE_SCALE_FACTOR * 2
+
+
+def run_transformer(
+    transformer,
+    latents: torch.Tensor,  # (B, C, h, w) noisy target latents
+    timestep: torch.Tensor,  # (B,) toolkit scale, 0..1000
+    query_embeds: List[torch.Tensor],  # per item (256, cap_feat_dim)
+    direct_embeds: List[torch.Tensor],  # per item (T, dim)
+    ref_latents: Optional[List[Optional[torch.Tensor]]] = None,  # per item (C, h, w) or None
+    **kwargs,
+) -> torch.Tensor:
+    """Returns the flow-matching velocity `noise - clean` as `(B, C, h, w)`."""
+    batch_size = latents.shape[0]
+    dtype = latents.dtype
+    x = list(latents.unsqueeze(2).unbind(dim=0))
+    ref_x = None
+    if ref_latents is not None and any(r is not None for r in ref_latents):
+        ref_x = [None if r is None else r.unsqueeze(1).to(dtype) for r in ref_latents]
+    tau = 1.0 - timestep.to(latents.device, dtype=torch.float32) / 1000.0
+    if tau.dim() == 0:
+        tau = tau.unsqueeze(0)
+    if tau.shape[0] != batch_size:
+        tau = tau.expand(batch_size)
+    out = transformer(
+        x,
+        tau.to(dtype),
+        [q.to(latents.device, dtype) for q in query_embeds],
+        cap_feats_2=[d.to(latents.device, dtype) for d in direct_embeds],
+        ref_x=ref_x,
+        return_dict=False,
+        **kwargs,
+    )[0]
+    # (C, 1, h, w) per item -> (B, C, h, w); the model predicts clean - noise
+    return -torch.stack([o[:, 0] for o in out], dim=0)
+
+
+class MingImagePipeline:
+    """Minimal flow-matching sampler for ai-toolkit's preview generation."""
+
+    def __init__(self, model):
+        # `model` is the MingImageModel, so its encode/decode and config are reused
+        self.model = model
+
+    @property
+    def device(self):
+        return self.model.device_torch
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def set_progress_bar_config(self, **kwargs):
+        pass
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        conditional_embeds,
+        unconditional_embeds=None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 12,
+        guidance_scale: float = 1.0,
+        latents: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+        reference_latents: Optional[List[Optional[torch.Tensor]]] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        model = self.model
+        device, dtype = model.device_torch, model.torch_dtype
+        transformer = model.transformer
+
+        latent_height = height // VAE_SCALE_FACTOR
+        latent_width = width // VAE_SCALE_FACTOR
+        channels = transformer.config.in_channels
+
+        if latents is None:
+            # randn_tensor, not torch.randn: the caller's generator may be a cpu one
+            latents = randn_tensor(
+                (1, channels, latent_height, latent_width),
+                generator=generator,
+                device=torch.device(device),
+                dtype=torch.float32,
+            )
+        latents = latents.to(device, dtype=dtype)
+
+        scheduler = model.get_train_scheduler()
+        # the reference pipeline spaces `steps` sigmas from 1 down to 0 and
+        # evaluates the model at sigma 0 too (a no-op step); skip that one
+        if num_inference_steps > 1:
+            sigmas = np.linspace(1.0, 0.0, num_inference_steps)[:-1]
+        else:
+            sigmas = np.array([1.0])
+        mu = ming_shift_mu((height // PIXELS_PER_TOKEN) * (width // PIXELS_PER_TOKEN))
+        scheduler.set_timesteps(sigmas=sigmas, device=device, mu=mu)
+        scheduler.set_begin_index(0)
+
+        query = list(conditional_embeds.text_embeds)
+        direct = list(conditional_embeds.direct_embeds)
+        # the reference negative is zero conditioning, whatever the negative prompt says
+        do_cfg = guidance_scale > 1.0
+        query_neg = [torch.zeros_like(q) for q in query]
+        direct_neg = [torch.zeros_like(d) for d in direct]
+
+        for timestep in scheduler.timesteps:
+            t = timestep.expand(latents.shape[0]).to(device)
+            noise_pred = run_transformer(
+                transformer, latents, t, query, direct, ref_latents=reference_latents
+            )
+            if do_cfg:
+                uncond_pred = run_transformer(
+                    transformer, latents, t, query_neg, direct_neg, ref_latents=reference_latents
+                )
+                noise_pred = uncond_pred + guidance_scale * (noise_pred - uncond_pred)
+
+            latents = scheduler.step(
+                noise_pred.to(torch.float32),
+                timestep,
+                latents.to(torch.float32),
+                return_dict=False,
+            )[0].to(dtype)
+
+        return model.decode_to_images(latents)

extensions_built_in/diffusion_models/ming_image/src/text_encoder.py

@@ -0,0 +1,706 @@
+"""Ming-Image conditioning stack: the multimodal LLM plus the projections that
+turn its hidden states into the two caption streams the DiT reads.
+
+The checkpoint spreads this over three folders; one module holds it all so the
+toolkit can quantize, offload and unload it as a single text encoder:
+
+  mllm/       Ling-mini-2.0 MoE decoder (`llm`) + Qwen2.5-VL vision tower
+              (`vision`) + its projector (`linear_proj`)
+  mlp/        the 256 learnable query tokens, `proj_in` / `proj_out` around the
+              connector, and `proj_directvlm`
+  connector/  a Qwen2-1.5B run bidirectionally over the query-token states
+
+Encoding a prompt:
+
+  system + human turn (+ reference block) + assistant prefix
+    + `<image>` + 256 x `<imagePatch>` + `</image>`
+    -> LLM (query slots carry the learnable tokens, routed via image_gate)
+    -> final states at the 256 query slots
+         -> proj_in -> connector -> proj_out          = query condition (256, 2560)
+    -> residual stream entering layers 5 and 12 plus the final normalized
+       output, concatenated, at every token BEFORE the query block
+         -> RMSNorm -> proj_directvlm                  = direct condition (T, 3840)
+
+A reference image (editing) is a `<image>...</image>` block ahead of the
+prompt text whose slots carry vision-tower features. Its tokens are part of the
+direct condition; the DiT also gets its VAE latent as a reference frame.
+
+The ComfyUI repack folds the same pieces, minus the vision tower, into one
+file (`convert_state_dict_on_load`); the tower then comes from the vendor
+checkpoint's shards.
+"""
+
+import json
+import os
+import random
+from collections import OrderedDict
+from typing import Callable, Dict, List, Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from accelerate import init_empty_weights
+from huggingface_hub import hf_hub_download, snapshot_download
+from PIL import Image
+from safetensors import safe_open
+from transformers import AutoTokenizer, Qwen2Config, Qwen2Model, Qwen2VLImageProcessor
+from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+    Qwen2_5_VisionTransformerPretrainedModel,
+)
+
+from toolkit.models.v2._mixin import OstrisTransformersMixin
+from toolkit.util.comfy_quant_import import parse_comfy_quant_blob
+
+from .bailing_moe_v2 import (
+    BailingMoeV2Config,
+    BailingMoeV2Model,
+    FusedExperts,
+    FusedExpertsMemoryManager,
+    RMSNorm,
+    get_video_rope_index,
+)
+from .checkpoints import BASE_REPO, COMFY_REPO, COMFY_TEXT_ENCODER_FILES, comfy_weight_names
+
+IMAGE_START = "<image>"
+IMAGE_PATCH = "<imagePatch>"
+IMAGE_END = "</image>"
+
+# the vendor processor's template (processing_bailingmm2.py), NOT the
+# tokenizer's own chat template: that one drops the assistant persona line
+SYSTEM_PROMPT = "<role>SYSTEM</role>你是一个友好的AI助手。\n\ndetailed thinking off"
+HUMAN_PREFIX = "<role>HUMAN</role>"
+ASSISTANT_PREFIX = "<role>ASSISTANT</role>"
+
+# subfolders of the checkpoint this module is assembled from
+MLLM_DIR = "mllm"
+MLP_DIR = "mlp"
+CONNECTOR_DIR = "connector"
+
+_VISION_CONFIG_KEYS = (
+    "depth",
+    "hidden_size",
+    "hidden_act",
+    "intermediate_size",
+    "num_heads",
+    "in_channels",
+    "patch_size",
+    "spatial_merge_size",
+    "temporal_patch_size",
+    "tokens_per_second",
+    "window_size",
+    "out_hidden_size",
+    "fullatt_block_indexes",
+)
+
+
+# the checkpoint's configs, tokenizer and image processor
+CONFIG_PATTERNS = [f"{MLLM_DIR}/*.json", f"{MLP_DIR}/config.json", f"{CONNECTOR_DIR}/config.json"]
+
+
+def resolve_checkpoint_root(name_or_path: str, weights: bool = True) -> str:
+    """The directory holding mllm/, mlp/ and connector/: a local checkpoint
+    root (or one of its component folders), or the hub snapshot with just
+    those folders fetched. `weights=False` fetches only the configs, tokenizer
+    and image processor (a ComfyUI single file supplies the weights)."""
+    if os.path.isdir(name_or_path):
+        root = os.path.abspath(name_or_path)
+        if os.path.basename(root) in (MLLM_DIR, MLP_DIR, CONNECTOR_DIR):
+            root = os.path.dirname(root)
+        return root
+    patterns = [f"{d}/*" for d in (MLLM_DIR, MLP_DIR, CONNECTOR_DIR)] if weights else CONFIG_PATTERNS
+    return snapshot_download(name_or_path, allow_patterns=patterns)
+
+
+def _safetensors_files(folder: str) -> List[str]:
+    files = sorted(f for f in os.listdir(folder) if f.endswith(".safetensors"))
+    if not files:
+        raise FileNotFoundError(f"no .safetensors files under {folder}")
+    return [os.path.join(folder, f) for f in files]
+
+
+class MingImageTextEncoder(nn.Module, OstrisTransformersMixin):
+    aitk_subfolder = MLLM_DIR
+    aitk_config_repo = BASE_REPO
+    aitk_comfy_repo = COMFY_REPO
+    aitk_comfy_weight_names = comfy_weight_names(COMFY_TEXT_ENCODER_FILES)
+
+    def __init__(
+        self,
+        llm_config: BailingMoeV2Config,
+        vision_config: Qwen2_5_VLVisionConfig,
+        connector_config: Qwen2Config,
+        mlp_config: dict,
+    ):
+        super().__init__()
+        hidden = llm_config.hidden_size
+        self.llm_config = llm_config
+        self.mlp_config = dict(mlp_config)
+        scales = list(mlp_config.get("img_gen_scales", [16]))
+        if scales != [16]:
+            raise ValueError(f"only img_gen_scales=[16] is supported, got {scales}")
+        self.num_query_tokens = 16 * 16
+        self.selected_layers = tuple(int(i) for i in mlp_config["selected_hidden_states_layers"])
+        if not mlp_config.get("use_identity_mlp", False):
+            raise NotImplementedError("Ming-Image with a non-identity diffusion MLP is not supported")
+        # L2-normalize the connector output / rescale like the original text encoder
+        self.normalize_query = bool(mlp_config.get("connector_norm", False))
+        self.text_encoder_norm = bool(mlp_config.get("text_encoder_norm", False))
+
+        self.llm = BailingMoeV2Model(llm_config)
+        self.vision = Qwen2_5_VisionTransformerPretrainedModel(vision_config)
+        self.linear_proj = nn.Sequential(
+            nn.Linear(vision_config.out_hidden_size, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+        )
+        self.query_tokens = nn.Parameter(torch.empty(self.num_query_tokens, hidden))
+
+        self.proj_in = nn.Linear(hidden, connector_config.hidden_size)
+        self.connector = Qwen2Model(connector_config)
+        # the connector only ever sees inputs_embeds (the projected query
+        # states): its token embedding is 467 MB of dead weight, and the
+        # ComfyUI repack omits it
+        self.connector.embed_tokens = None
+        self.proj_out = nn.Linear(connector_config.hidden_size, int(mlp_config["diffusion_c_input_dim"]))
+
+        direct_dim = hidden * len(self.selected_layers)
+        self.proj_directvlm = nn.Sequential(
+            RMSNorm(direct_dim, eps=1e-5),
+            nn.Linear(direct_dim, int(mlp_config["diffusion_inner_dim"])),
+        )
+
+    # ------------------------------------------------------------------
+    # toolkit hooks
+    # ------------------------------------------------------------------
+    @property
+    def device(self):
+        return self.query_tokens.device
+
+    @property
+    def dtype(self):
+        return self.query_tokens.dtype
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["llm.layers", "connector.layers", "vision.blocks"]
+
+    def _apply(self, fn, *args, **kwargs):
+        # a device move or dtype cast invalidates captured CUDA graphs (weight
+        # addresses are baked in); the prompt encoder registers its caches
+        # here so they release their device memory before the move
+        for callback in getattr(self, "_on_apply", ()):
+            callback()
+        return super()._apply(fn, *args, **kwargs)
+
+    @classmethod
+    def get_quantization_exclude_modules(cls):
+        return ["query_tokens*", "proj_in*", "proj_out*", "proj_directvlm*", "linear_proj*", "llm.word_embeddings*"]
+
+    # ------------------------------------------------------------------
+    # loading
+    # ------------------------------------------------------------------
+    @classmethod
+    def build_configs(cls, root: str):
+        with open(os.path.join(root, MLLM_DIR, "config.json")) as f:
+            mllm = json.load(f)
+        with open(os.path.join(root, MLP_DIR, "config.json")) as f:
+            mlp = json.load(f)
+        llm_config = BailingMoeV2Config(**mllm["llm_config"])
+        vision_dict = {k: v for k, v in mllm["vision_config"].items() if k in _VISION_CONFIG_KEYS}
+        vision_config = Qwen2_5_VLVisionConfig(**vision_dict)
+        vision_config._attn_implementation = "sdpa"
+        connector_config = Qwen2Config.from_pretrained(os.path.join(root, CONNECTOR_DIR))
+        connector_config._attn_implementation = "sdpa"
+        connector_config.use_cache = False
+        return llm_config, vision_config, connector_config, mlp
+
+    @classmethod
+    def aitk_from_pretrained(cls, path, subfolder=None, dtype=None, **kwargs):
+        dtype = dtype or torch.bfloat16
+        root = resolve_checkpoint_root(path)
+        llm_config, vision_config, connector_config, mlp = cls.build_configs(root)
+        with init_empty_weights(include_buffers=False):
+            model = cls(llm_config, vision_config, connector_config, mlp)
+
+        state = {}
+
+        def read(folder: str, rename):
+            for file in _safetensors_files(folder):
+                with safe_open(file, framework="pt") as f:
+                    for key in f.keys():
+                        new_key = rename(key)
+                        if new_key is None:
+                            continue
+                        value = f.get_tensor(key)
+                        if value.is_floating_point():
+                            value = value.to(dtype)
+                        state[new_key] = value
+
+        def rename_mllm(key: str):
+            if key.startswith("model.model."):
+                key = "llm." + key[len("model.model."):]
+                return None if ".mlp.audio_gate." in key else key
+            if key.startswith("vision.") or key.startswith("linear_proj."):
+                return key
+            return None  # lm_head, audio pieces
+
+        def rename_mlp(key: str):
+            if key == "query_tokens_dict.16x16":
+                return "query_tokens"
+            if key.startswith(("proj_in.", "proj_out.", "proj_directvlm.")):
+                return key
+            return None
+
+        def rename_connector(key: str):
+            if not key.startswith("model.") or key.startswith("model.embed_tokens."):
+                return None
+            return "connector." + key[len("model."):]
+
+        read(os.path.join(root, MLLM_DIR), rename_mllm)
+        read(os.path.join(root, MLP_DIR), rename_mlp)
+        read(os.path.join(root, CONNECTOR_DIR), rename_connector)
+
+        # the checkpoint stores one linear per expert; stack them into the
+        # grouped-GEMM banks (the per-expert tensors are dropped as they go)
+        for i, layer in enumerate(model.llm.layers):
+            if layer.is_moe:
+                FusedExperts.fuse_state_dict(state, f"llm.layers.{i}.mlp.experts.", llm_config.num_experts)
+
+        result = model.load_state_dict(state, assign=True, strict=False)
+        if result.missing_keys or result.unexpected_keys:
+            raise ValueError(
+                f"MingImageTextEncoder load mismatch: missing {result.missing_keys[:8]}, "
+                f"unexpected {result.unexpected_keys[:8]}"
+            )
+        leftover = [n for n, p in model.named_parameters() if p.is_meta]
+        if leftover:
+            raise ValueError(f"MingImageTextEncoder load left meta parameters: {leftover[:8]}")
+        model.to(dtype=dtype)
+        model.eval()
+        return model
+
+    # ------------------------------------------------------------------
+    # single-file (ComfyUI repack) loading
+    # ------------------------------------------------------------------
+    @classmethod
+    def aitk_load_config(cls, path, subfolder=None):
+        return cls.build_configs(resolve_checkpoint_root(path, weights=False))
+
+    @classmethod
+    def aitk_from_config(cls, config):
+        with init_empty_weights(include_buffers=False):
+            return cls(*config)
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        """The ComfyUI repack's keys -> this module's. The LLM sits under
+        `thinker.` with transformers-style names and fused expert banks in our
+        `(E, 2*inter, hidden)` / `(E, hidden, inter)` layout; the connector,
+        projections and query tokens already match. Quantized expert banks
+        (int8 rows, float32 scales, optionally convrot-rotated) become the
+        `*_q` / `*_scale` entries FusedExperts attaches on load; the toolkit's
+        importer covers the nn.Linear markers."""
+        if not any(k.startswith("thinker.") for k in state_dict):
+            return state_dict
+        renames = (
+            (".mlp.gate.proj.", ".mlp.gate."),
+            (".mlp.image_gate.proj.", ".mlp.image_gate."),
+        )
+        out = {}
+        banks = {}
+        for key, value in state_dict.items():
+            if key == "tokenizer_json":  # the tokenizer comes from the vendor repo
+                continue
+            if not key.startswith("thinker."):
+                out[key] = value
+                continue
+            key = "llm." + key[len("thinker."):]
+            if key.startswith("llm.embed_tokens."):
+                key = "llm.word_embeddings." + key[len("llm.embed_tokens."):]
+            for old, new in renames:
+                key = key.replace(old, new)
+            if ".mlp.experts." in key:
+                prefix, _, rest = key.partition(".mlp.experts.")
+                bank, _, field = rest.partition(".")  # gate_up_proj.<field>
+                banks.setdefault(prefix + ".mlp.experts.", {})[(bank[: -len("_proj")], field)] = value
+                continue
+            out[key] = value
+        for prefix, entries in banks.items():
+            if not any(field == "comfy_quant" for _, field in entries):
+                for (bank, field), value in entries.items():
+                    if field != "weight":
+                        raise ValueError(f"{prefix}{bank}: unexpected entry {field!r}")
+                    out[f"{prefix}{bank}"] = value
+                continue
+            rot = None
+            for bank in ("gate_up", "down"):
+                conf = parse_comfy_quant_blob(entries[(bank, "comfy_quant")])
+                if conf.get("format") != "int8_tensorwise":
+                    raise ValueError(
+                        f"{prefix}{bank}: unsupported expert quantization {conf.get('format')!r}"
+                    )
+                bank_rot = int(conf.get("convrot_groupsize", 256)) if conf.get("convrot") else 1
+                if rot is not None and bank_rot != rot:
+                    raise ValueError(f"{prefix}: the two expert banks use different rotations")
+                rot = bank_rot
+                out[f"{prefix}{bank}_q"] = entries[(bank, "weight")]
+                out[f"{prefix}{bank}_scale"] = entries[(bank, "weight_scale")]
+            out[f"{prefix}rot_size"] = torch.tensor(rot)
+        return out
+
+    @classmethod
+    def load_from_state_dict(
+        cls, state_dict, dtype, config_path=None, config=None, subfolder=None, **kwargs
+    ):
+        state_dict = cls.convert_state_dict_on_load(state_dict)
+        if not any(k.startswith("vision.") for k in state_dict):
+            # the ComfyUI repack carries no vision tower (its PR is text to
+            # image so far); editing needs one, so it comes from the vendor repo
+            state_dict = dict(state_dict)
+            state_dict.update(cls.read_vision_weights(config_path or cls.aitk_config_repo, dtype))
+        return super().load_from_state_dict(
+            state_dict, dtype, config_path=config_path, config=config, subfolder=subfolder, **kwargs
+        )
+
+    @classmethod
+    def read_vision_weights(cls, name_or_path: str, dtype) -> Dict[str, torch.Tensor]:
+        """`vision.*` and `linear_proj.*` from the vendor checkpoint's mllm
+        shards: only the shards holding those keys are fetched, and only
+        those tensors are read."""
+
+        def wanted(key: str) -> bool:
+            return key.startswith(("vision.", "linear_proj."))
+
+        root = resolve_checkpoint_root(name_or_path, weights=False)
+        folder = os.path.join(root, MLLM_DIR)
+        index_path = os.path.join(folder, "model.safetensors.index.json")
+        if os.path.isfile(index_path):
+            with open(index_path) as f:
+                weight_map = json.load(f)["weight_map"]
+            shards = sorted({v for k, v in weight_map.items() if wanted(k)})
+        else:
+            shards = [os.path.basename(f) for f in _safetensors_files(folder)]
+        state = {}
+        for shard in shards:
+            path = os.path.join(folder, shard)
+            if not os.path.isfile(path):
+                if os.path.isdir(name_or_path):
+                    raise FileNotFoundError(f"{path} is missing")
+                path = hf_hub_download(name_or_path, f"{MLLM_DIR}/{shard}")
+            with safe_open(path, framework="pt") as f:
+                for key in f.keys():
+                    if wanted(key):
+                        state[key] = f.get_tensor(key).to(dtype)
+        if not state:
+            raise ValueError(f"no vision tower weights under {folder}")
+        return state
+
+    def aitk_post_load(
+        self, qtype=None, offload=0.0, quantize_device=None, device=None, dtype=torch.bfloat16, **kwargs
+    ):
+        # the expert banks are not nn.Linear, so the toolkit's quantizer,
+        # dequantizer and layer offloader never see them. They follow the
+        # request in kind before the regular pass handles the attention /
+        # dense / shared-expert linears with `qtype`: any quantize_te makes
+        # them int8 weight-only (ready-made comfy int8 banks stay as they are),
+        # none restores full precision. With layer offloading on, the banks get
+        # their own stager (pinned cpu, copied in per forward) so the offloader
+        # leaves them where they are instead of making them resident.
+        q_device = quantize_device if quantize_device is not None else device
+        banks = [layer.mlp.experts for layer in self.llm.layers if layer.is_moe]
+        for bank in banks:
+            if qtype:
+                bank.quantize_int8_(device=q_device)
+            else:
+                bank.dequantize_(dtype=dtype, device=q_device)
+        if offload and offload > 0 and q_device is not None:
+            for bank in banks:
+                if random.random() <= offload:
+                    FusedExpertsMemoryManager.attach(bank, torch.device(q_device))
+        return super().aitk_post_load(
+            qtype=qtype, offload=offload, quantize_device=quantize_device, device=device, dtype=dtype, **kwargs
+        )
+
+    @classmethod
+    def load_tokenizer_and_processor(cls, name_or_path: str):
+        root = resolve_checkpoint_root(name_or_path, weights=False)
+        folder = os.path.join(root, MLLM_DIR)
+        tokenizer = AutoTokenizer.from_pretrained(folder)
+        image_processor = Qwen2VLImageProcessor.from_pretrained(folder)
+        return tokenizer, image_processor
+
+    # ------------------------------------------------------------------
+    # forward pieces
+    # ------------------------------------------------------------------
+    def encode_reference_images(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
+        """Vision-tower features for the reference blocks, `(N, hidden)` in
+        token order across all images."""
+        features = self.vision(pixel_values.to(self.dtype), grid_thw=grid_thw).pooler_output
+        return F.normalize(self.linear_proj(features), dim=-1)
+
+    def encode_query(self, query_states: torch.Tensor) -> torch.Tensor:
+        """`(B, 256, hidden)` final LLM states at the query slots -> `(B, 256, c_dim)`."""
+        x = self.proj_in(query_states)
+        bsz, n, _ = x.shape
+        # the connector is a decoder run as a bidirectional encoder
+        mask = torch.ones((bsz, 1, n, n), dtype=torch.bool, device=x.device)
+        hidden = self.connector(inputs_embeds=x, attention_mask=mask, use_cache=False).last_hidden_state
+        out = self.proj_out(hidden)
+        if self.normalize_query:
+            out = F.normalize(out, dim=-1)
+        if self.text_encoder_norm:
+            out = out * 1000.0
+        return out
+
+    def encode_direct(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """`(T, hidden * n_selected)` concatenated layer states -> `(T, dim)`."""
+        return self.proj_directvlm(hidden_states)
+
+
+# prompts are padded up to a multiple of this many tokens so a few captured
+# graphs cover every length; longer sequences run eagerly
+GRAPH_LENGTH_STEP = 64
+GRAPH_MAX_LENGTH = 2048
+GRAPH_MAX_ENTRIES = 8
+
+
+class CudaGraphCache:
+    """Replays a function through CUDA graphs, one capture per static-shape key.
+
+    Encoding a prompt is ~3k small launches on ~30 ms of GPU work, so at
+    batch 1 the encoder is CPU bound; a captured graph replays the whole
+    stage in one call. Inputs are copied into the captured tensors before
+    each replay; outputs are the captured tensors themselves (overwritten by
+    the next replay, so callers copy what they keep). Captures are dropped
+    when the module's weights move (their addresses are baked into the graph).
+    """
+
+    def __init__(self, fn: Callable, signature: Callable[[], tuple], max_entries: int = GRAPH_MAX_ENTRIES):
+        self.fn = fn
+        self.signature = signature
+        self.max_entries = max_entries
+        self.entries: "OrderedDict[tuple, tuple]" = OrderedDict()
+        self.captured_signature = None
+
+    def clear(self):
+        self.entries.clear()
+        self.captured_signature = None
+
+    def __call__(self, key: tuple, inputs: Dict[str, torch.Tensor], **static_kwargs):
+        signature = self.signature()
+        if signature != self.captured_signature:
+            self.clear()
+            self.captured_signature = signature
+        entry = self.entries.get(key)
+        if entry is None:
+            if len(self.entries) >= self.max_entries:
+                self.entries.popitem(last=False)
+            static = {name: value.clone() for name, value in inputs.items()}
+            stream = torch.cuda.Stream(device=next(iter(static.values())).device)
+            stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(stream):
+                # warm-up on a side stream: lazy allocations / kernel compiles
+                # must not land inside the capture
+                for _ in range(2):
+                    self.fn(**static, **static_kwargs)
+            torch.cuda.current_stream().wait_stream(stream)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                outputs = self.fn(**static, **static_kwargs)
+            entry = (graph, static, outputs)
+            self.entries[key] = entry
+        else:
+            self.entries.move_to_end(key)
+        graph, static, outputs = entry
+        for name, value in inputs.items():
+            static[name].copy_(value)
+        graph.replay()
+        return outputs
+
+
+class MingImagePromptEncoder:
+    """Prompts (+ optional single reference image each) -> the DiT's two
+    caption streams, one entry per prompt.
+
+    On a resident GPU text encoder the LLM (per padded-length bucket) and the
+    connector (fixed 256 tokens) replay as CUDA graphs; set `use_cuda_graphs`
+    False for eager execution. Offloaded or cpu encoders always run eagerly.
+    """
+
+    use_cuda_graphs = True
+
+    def __init__(self, text_encoder: MingImageTextEncoder, tokenizer, image_processor):
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.image_processor = image_processor
+        self.merge_size = int(getattr(image_processor, "merge_size", 2))
+        self.patch_token_id = tokenizer.convert_tokens_to_ids(IMAGE_PATCH)
+        self.start_token_id = tokenizer.convert_tokens_to_ids(IMAGE_START)
+        self.end_token_id = tokenizer.convert_tokens_to_ids(IMAGE_END)
+        self.pad_token_id = tokenizer.pad_token_id
+        config = text_encoder.llm_config
+        te = text_encoder
+        self._llm_graphs = CudaGraphCache(te.llm, self._weight_signature)
+        self._connector_graphs = CudaGraphCache(te.encode_query, self._weight_signature)
+        te._on_apply = [self._llm_graphs.clear, self._connector_graphs.clear]
+        if (self.patch_token_id, self.start_token_id, self.end_token_id) != (
+            config.image_patch_token,
+            config.image_start_token,
+            config.image_end_token,
+        ):
+            raise ValueError("tokenizer image tokens do not match the LLM config")
+
+    def _query_block(self) -> str:
+        return IMAGE_START + IMAGE_PATCH * self.text_encoder.num_query_tokens + IMAGE_END
+
+    def _chat_text(self, content: str) -> str:
+        eos = self.tokenizer.eos_token
+        return SYSTEM_PROMPT + eos + HUMAN_PREFIX + content + eos + ASSISTANT_PREFIX
+
+    def _weight_signature(self) -> tuple:
+        # a handful of storages stands in for "the weights moved": any device
+        # move relocates all of them
+        te = self.text_encoder
+        probes = [te.query_tokens, te.proj_out.weight]
+        for module in (te.llm.layers[0], te.llm.layers[-1], te.connector.layers[0], te.connector.layers[-1]):
+            probe = next(module.parameters(), None)
+            probes.append(probe if probe is not None else next(module.buffers()))
+        return tuple(int(t.data_ptr()) for t in probes)
+
+    def _graphs_enabled(self) -> bool:
+        te = self.text_encoder
+        if not self.use_cuda_graphs or te.device.type != "cuda":
+            return False
+        # streamed weights change addresses every forward
+        return not hasattr(te, "_memory_manager") and not any(
+            hasattr(m, "_layer_memory_manager") for m in te.modules()
+        )
+
+    @torch.no_grad()
+    def encode(
+        self,
+        prompts: Sequence[str],
+        images: Optional[Sequence[Sequence[Image.Image]]] = None,
+    ) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[bool]]:
+        te = self.text_encoder
+        device, dtype = te.device, te.dtype
+        images = list(images) if images is not None else [[] for _ in prompts]
+
+        texts, pixel_values, grids, ref_counts = [], [], [], []
+        for prompt, sample_images in zip(prompts, images):
+            if len(sample_images) > 1:
+                raise ValueError("Ming-Image takes at most one reference image per prompt")
+            blocks = ""
+            for image in sample_images:
+                feats = self.image_processor(images=[image], return_tensors="pt")
+                grid = feats["image_grid_thw"][0]
+                n_tokens = int(grid.prod()) // (self.merge_size**2)
+                blocks += IMAGE_START + IMAGE_PATCH * n_tokens + IMAGE_END + "\n"
+                pixel_values.append(feats["pixel_values"])
+                grids.append(grid)
+            ref_counts.append(len(sample_images))
+            texts.append(self._chat_text(blocks + prompt) + self._query_block())
+
+        # the template supplies every special token; right padding keeps the
+        # query block at the end of each valid sequence
+        enc = self.tokenizer(
+            texts, add_special_tokens=False, padding=True, padding_side="right", return_tensors="pt"
+        )
+        input_ids = enc.input_ids
+        attention_mask = enc.attention_mask
+        use_graphs = self._graphs_enabled()
+        if use_graphs:
+            # pad to the bucket length; the pad slots are masked keys whose
+            # own outputs are never read, so the valid tokens are unaffected
+            padded = -(-input_ids.shape[1] // GRAPH_LENGTH_STEP) * GRAPH_LENGTH_STEP
+            if padded > GRAPH_MAX_LENGTH:
+                use_graphs = False
+            elif padded > input_ids.shape[1]:
+                extra = padded - input_ids.shape[1]
+                input_ids = F.pad(input_ids, (0, extra), value=self.pad_token_id)
+                attention_mask = F.pad(attention_mask, (0, extra), value=0)
+        input_ids = input_ids.to(device)
+        attention_mask = attention_mask.to(device)
+        bsz, seq_len = input_ids.shape
+
+        embeds = te.llm.word_embeddings(input_ids)
+        ref_feats = None
+        if pixel_values:
+            ref_feats = te.encode_reference_images(
+                torch.cat(pixel_values, dim=0).to(device), torch.stack(grids).to(device)
+            )
+
+        # fill the slots: reference features first (in order), then the query tokens
+        patch_mask = input_ids == self.patch_token_id
+        query_tokens = te.query_tokens.to(dtype)
+        ref_offset = 0
+        image_grid_thw = []
+        for b in range(bsz):
+            n_ref = 0
+            if ref_counts[b]:
+                grid = grids[ref_offset]
+                n_ref = int(grid.prod()) // (self.merge_size**2)
+                image_grid_thw.append(grid.tolist())
+            image_grid_thw.append([1, 2, self.text_encoder.num_query_tokens * 2])
+            fill = query_tokens
+            if n_ref:
+                fill = torch.cat([ref_feats[:n_ref].to(dtype), fill], dim=0)
+                ref_feats = ref_feats[n_ref:]
+                ref_offset += 1
+            positions = torch.nonzero(patch_mask[b], as_tuple=False).squeeze(1)
+            if positions.numel() != fill.shape[0]:
+                raise ValueError(
+                    f"prompt {b} has {positions.numel()} image slots but {fill.shape[0]} features"
+                )
+            embeds[b, positions] = fill
+
+        position_ids = get_video_rope_index(
+            te.llm_config,
+            input_ids,
+            attention_mask,
+            torch.tensor(image_grid_thw, dtype=torch.long),
+        ).to(device)
+        causal = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))
+        attn_mask = causal[None, None] & attention_mask.bool()[:, None, None, :]
+
+        capture = tuple(sorted(set(te.selected_layers) | {len(te.llm.layers)}))
+        llm_inputs = dict(
+            inputs_embeds=embeds.to(dtype),
+            attention_mask=attn_mask,
+            position_ids=position_ids,
+            image_mask=patch_mask,
+        )
+        if use_graphs:
+            final, captured = self._llm_graphs((bsz, seq_len), llm_inputs, capture_layers=capture)
+        else:
+            final, captured = te.llm(**llm_inputs, capture_layers=capture)
+
+        n_query = te.num_query_tokens
+        query_states, direct_inputs, direct_lens = [], [], []
+        for b in range(bsz):
+            length = int(attention_mask[b].sum())
+            tail = input_ids[b, length - n_query - 2 : length]
+            if int(tail[0]) != self.start_token_id or int(tail[-1]) != self.end_token_id:
+                raise ValueError("prompt does not end with the query block")
+            query_states.append(final[b, length - n_query - 1 : length - 1])
+            # the direct condition covers every token before the query block
+            direct_end = length - n_query - 2
+            direct_inputs.append(
+                torch.cat([captured[i][b, :direct_end] for i in te.selected_layers], dim=-1)
+            )
+            direct_lens.append(direct_end)
+
+        query_batch = torch.stack(query_states, dim=0)
+        if use_graphs:
+            query_cond = self._connector_graphs((bsz,), dict(query_states=query_batch))
+        else:
+            query_cond = te.encode_query(query_batch)
+        direct_batch = torch.nn.utils.rnn.pad_sequence(direct_inputs, batch_first=True)
+        direct_cond = te.encode_direct(direct_batch)
+
+        # clones: graph outputs are overwritten by the next replay
+        query_out = [query_cond[b].clone() for b in range(bsz)]
+        direct_out = [direct_cond[b, : direct_lens[b]].clone() for b in range(bsz)]
+        has_reference = [count > 0 for count in ref_counts]
+        return query_out, direct_out, has_reference

extensions_built_in/diffusion_models/ming_image/src/transformer.py

@@ -0,0 +1,591 @@
+# Copyright 2025 Alibaba Z-Image 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.
+
+"""Ming-Image's diffusion transformer: the Z-Image 6B DiT with Ming's changes.
+
+Vendored from diffusers' `transformer_z_image.py` (Apache 2.0) because the
+Ming checkpoint (`_class_name: DiffusionTransformer`) diverges from Z-Image in
+ways the diffusers class cannot express:
+
+  - `alignment_padding_mode="zero_masked"`: the reference pads sequences to a
+    multiple of 32 with zeros that are masked out of every attention and
+    dropped from the output. Those slots cannot influence anything, so this
+    port never materializes them: only their share of the position grid
+    survives (the image block starts after the padded caption length). With
+    nothing to mask, equal-length batches run unmasked (flash) attention.
+    There are no learned `x_pad_token` / `cap_pad_token`.
+  - a second caption stream, `cap_feats_2`: the direct-VLM condition already
+    lives in model width, so it bypasses `cap_embedder` and is appended to the
+    embedded query-token captions before the context refiner.
+  - `ref_x`: a reference latent concatenated along the frame axis (F=2) for
+    editing; with `multi_frame_output=False` only the first frame is returned.
+
+`OstrisModelMixin` gives it the toolkit's universal load/quantize/offload
+path. The omni/siglip and controlnet branches of the upstream file are gone.
+"""
+
+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_dispatch import dispatch_attention_fn
+from diffusers.models.attention_processor import Attention
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+
+from toolkit.models.v2._mixin import OstrisModelMixin
+from toolkit.models.v2.diffusion_models.z_image import ZImageTransformer2DModel
+
+from .checkpoints import BASE_REPO, COMFY_REPO, COMFY_TRANSFORMER_FILES, comfy_weight_names
+
+ADALN_EMBED_DIM = 256
+SEQ_MULTI_OF = 32
+
+LEARNED_PADDING = "learned"
+ZERO_MASKED_PADDING = "zero_masked"
+
+
+class RMSNorm(nn.Module):
+    """Same parameters as diffusers' RMSNorm, on torch's fused kernel."""
+
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps)
+
+
+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.autocast(device_type=t.device.type, 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
+        compute_dtype = getattr(self.mlp[0], "compute_dtype", None)
+        if weight_dtype.is_floating_point:
+            t_freq = t_freq.to(weight_dtype)
+        elif compute_dtype is not None:
+            t_freq = t_freq.to(compute_dtype)
+        return self.mlp(t_freq)
+
+
+class ZSingleStreamAttnProcessor:
+    """Single-stream attention with complex RoPE, on the diffusers Attention class."""
+
+    _attention_backend = None
+    _parallel_config = None
+
+    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))
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        def apply_rotary_emb(x_in: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+            with torch.autocast(device_type=x_in.device.type, enabled=False):
+                x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2))
+                freqs = freqs.unsqueeze(2)
+                x_out = torch.view_as_real(x * freqs).flatten(3)
+                return x_out.type_as(x_in)
+
+        if freqs_cis is not None:
+            query = apply_rotary_emb(query, freqs_cis)
+            key = apply_rotary_emb(key, freqs_cis)
+
+        dtype = query.dtype
+        query, key = query.to(dtype), key.to(dtype)
+
+        # (batch, seq_len) -> (batch, 1, 1, seq_len), broadcast over heads and queries
+        if attention_mask is not None and attention_mask.ndim == 2:
+            attention_mask = attention_mask[:, None, None, :]
+
+        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,
+        )
+        hidden_states = hidden_states.flatten(2, 3).to(dtype)
+
+        output = attn.to_out[0](hidden_states)
+        if len(attn.to_out) > 1:
+            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(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+@maybe_allow_in_graph
+class ZImageTransformerBlock(nn.Module):
+    def __init__(self, layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=True):
+        super().__init__()
+        self.dim = dim
+        self.head_dim = dim // n_heads
+        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=ZSingleStreamAttnProcessor(),
+        )
+        self.attention.norm_q = RMSNorm(self.head_dim, eps=1e-5) if qk_norm else None
+        self.attention.norm_k = RMSNorm(self.head_dim, eps=1e-5) if qk_norm else None
+        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, attn_mask, freqs_cis, adaln_input=None):
+        if self.modulation:
+            mod = self.adaLN_modulation(adaln_input)
+            scale_msa, gate_msa, scale_mlp, gate_mlp = mod.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
+
+            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)
+            x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))
+        else:
+            attn_out = self.attention(self.attention_norm1(x), attention_mask=attn_mask, freqs_cis=freqs_cis)
+            x = x + self.attention_norm2(attn_out)
+            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)
+        return self.linear(x)
+
+
+class RopeEmbedder:
+    def __init__(self, theta=256.0, axes_dims=(16, 56, 56), axes_lens=(64, 128, 128)):
+        self.theta = theta
+        self.axes_dims = axes_dims
+        self.axes_lens = axes_lens
+        assert len(axes_dims) == len(axes_lens)
+        self.freqs_cis = None
+
+    @staticmethod
+    def precompute_freqs_cis(dim, end, theta=256.0):
+        with torch.device("cpu"):
+            freqs_cis = []
+            for d, e in 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.append(torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64))
+            return freqs_cis
+
+    def __call__(self, ids: torch.Tensor):
+        assert ids.ndim == 2 and 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 = [f.to(device) for f in self.freqs_cis]
+        elif self.freqs_cis[0].device != device:
+            self.freqs_cis = [f.to(device) for f in self.freqs_cis]
+        return torch.cat([self.freqs_cis[i][ids[:, i]] for i in range(len(self.axes_dims))], dim=-1)
+
+
+class MingImageTransformer2DModel(
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin
+):
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["ZImageTransformerBlock"]
+    _repeated_blocks = ["ZImageTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["t_embedder", "cap_embedder"]
+
+    aitk_subfolder = "transformer"
+    aitk_config_repo = BASE_REPO
+    aitk_comfy_repo = COMFY_REPO
+    aitk_comfy_weight_names = comfy_weight_names(COMFY_TRANSFORMER_FILES)
+
+    @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,
+        siglip_feat_dim=None,  # accepted for config compatibility, unused
+        rope_theta=256.0,
+        t_scale=1000.0,
+        axes_dims=(32, 48, 48),
+        axes_lens=(20480, 512, 512),
+        alignment_padding_mode=ZERO_MASKED_PADDING,
+        multi_frame_output=False,
+        # the trainer's dynamic-shift schedule reads config.patch_size to count
+        # image tokens; latents are patchified 2x2 (all_patch_size[0])
+        patch_size=2,
+    ) -> None:
+        super().__init__()
+        if alignment_padding_mode not in (LEARNED_PADDING, ZERO_MASKED_PADDING):
+            raise ValueError(f"unsupported alignment_padding_mode {alignment_padding_mode!r}")
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.all_patch_size = tuple(all_patch_size)
+        self.all_f_patch_size = tuple(all_f_patch_size)
+        self.dim = dim
+        self.n_heads = n_heads
+        self.rope_theta = rope_theta
+        self.t_scale = t_scale
+        self.alignment_padding_mode = alignment_padding_mode
+        self.multi_frame_output = multi_frame_output
+        self.gradient_checkpointing = False
+
+        assert len(all_patch_size) == len(all_f_patch_size)
+
+        all_x_embedder, all_final_layer = {}, {}
+        for p, fp in zip(all_patch_size, all_f_patch_size):
+            all_x_embedder[f"{p}-{fp}"] = nn.Linear(fp * p * p * in_channels, dim, bias=True)
+            all_final_layer[f"{p}-{fp}"] = FinalLayer(dim, p * p * fp * self.out_channels)
+        self.all_x_embedder = nn.ModuleDict(all_x_embedder)
+        self.all_final_layer = nn.ModuleDict(all_final_layer)
+
+        self.noise_refiner = nn.ModuleList(
+            [
+                ZImageTransformerBlock(1000 + i, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=True)
+                for i in range(n_refiner_layers)
+            ]
+        )
+        self.context_refiner = nn.ModuleList(
+            [
+                ZImageTransformerBlock(i, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation=False)
+                for i 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 alignment_padding_mode == LEARNED_PADDING:
+            self.x_pad_token = nn.Parameter(torch.zeros((1, dim)))
+            self.cap_pad_token = nn.Parameter(torch.zeros((1, dim)))
+        else:
+            self.register_parameter("x_pad_token", None)
+            self.register_parameter("cap_pad_token", None)
+
+        self.layers = nn.ModuleList(
+            [ZImageTransformerBlock(i, dim, n_heads, n_kv_heads, norm_eps, qk_norm) for i in range(n_layers)]
+        )
+        assert dim // n_heads == sum(axes_dims)
+        self.axes_dims = list(axes_dims)
+        self.axes_lens = list(axes_lens)
+        self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=self.axes_dims, axes_lens=self.axes_lens)
+
+    # ------------------------------------------------------------------
+    # toolkit hooks
+    # ------------------------------------------------------------------
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]
+
+    def get_offload_ignore_modules(self):
+        return [p for p in (self.x_pad_token, self.cap_pad_token) if p is not None]
+
+    @classmethod
+    def get_quantization_exclude_modules(cls):
+        # same precision-sensitive set as Z-Image: the timestep embedder feeds
+        # every adaLN, the caption / latent projections and the final layers
+        return ["t_embedder*", "cap_embedder*", "all_x_embedder*", "all_final_layer*"]
+
+    # ------------------------------------------------------------------
+    # single-file layout: the ComfyUI repack's int8 file uses Comfy's Z-Image
+    # names (fused attention.qkv, x_embedder / final_layer at the root); its
+    # bf16 file and the vendor's are diffusers-keyed and pass through. Saves
+    # use the Comfy names so ComfyUI and this class both load the result.
+    # ------------------------------------------------------------------
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        return ZImageTransformer2DModel.convert_state_dict_on_load(state_dict)
+
+    @classmethod
+    def convert_state_dict_on_save(cls, state_dict):
+        return ZImageTransformer2DModel.convert_state_dict_on_save(state_dict)
+
+    # ------------------------------------------------------------------
+    # sequence assembly
+    # ------------------------------------------------------------------
+    @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)]
+        return torch.stack(torch.meshgrid(axes, indexing="ij"), dim=-1)
+
+    def _positions(self, ori_len: int, pad_len: int, grid_size, grid_start, device):
+        """Position ids for one sequence: the grid, then `pad_len` pad slots
+        at the origin (learned padding only)."""
+        pos_ids = self.create_coordinate_grid(size=grid_size, start=grid_start, device=device).flatten(0, 2)
+        if pad_len > 0:
+            pad_pos = self.create_coordinate_grid(size=(1, 1, 1), start=(0, 0, 0), device=device)
+            pos_ids = torch.cat([pos_ids, pad_pos.flatten(0, 2).repeat(pad_len, 1)], dim=0)
+        return pos_ids
+
+    def _pad_feat(self, feat: torch.Tensor, pad_len: int) -> torch.Tensor:
+        # learned padding repeats the last feature; the pad token replaces it
+        if pad_len == 0:
+            return feat
+        return torch.cat([feat, feat[-1:].repeat(pad_len, 1)], dim=0)
+
+    def _apply_pad_token(self, feats: torch.Tensor, pad_lens: List[int], seqlens: List[int], pad_token):
+        if pad_token is None or not any(pad_lens):
+            return feats
+        pad_mask = torch.cat(
+            [
+                torch.arange(n, device=feats.device) >= n - pad
+                for n, pad in zip(seqlens, pad_lens)
+            ]
+        )
+        fill = pad_token.expand(feats.shape[0], -1).to(feats.dtype)
+        return torch.where(pad_mask.unsqueeze(1).expand_as(feats), fill, feats)
+
+    @staticmethod
+    def _length_mask(seqlens: List[int], device) -> Optional[torch.Tensor]:
+        """Key mask for right-padded sequences, None when nothing is padded
+        (decided on the host: no device sync)."""
+        max_len = max(seqlens)
+        if all(n == max_len for n in seqlens):
+            return None
+        mask = torch.zeros((len(seqlens), max_len), dtype=torch.bool, device=device)
+        for i, n in enumerate(seqlens):
+            mask[i, :n] = True
+        return mask
+
+    def _batch(self, feats: List[torch.Tensor], pos_ids: List[torch.Tensor], device):
+        """Right-pad per-item sequences into a batch; returns the batch, its
+        rope, a key mask (None when lengths agree) and the item lengths."""
+        seqlens = [len(f) for f in feats]
+        freqs = list(self.rope_embedder(torch.cat(pos_ids, dim=0)).split([len(p) for p in pos_ids], dim=0))
+        feats = pad_sequence(feats, batch_first=True, padding_value=0.0)
+        freqs = pad_sequence(freqs, batch_first=True, padding_value=0.0)[:, : feats.shape[1]]
+        return feats, freqs, self._length_mask(seqlens, device), seqlens
+
+    def unpatchify(self, x: List[torch.Tensor], sizes: List[Tuple[int, int, int]], patch_size, f_patch_size):
+        pH = pW = patch_size
+        pF = f_patch_size
+        out = []
+        for item, (frames, height, width) in zip(x, sizes):
+            ori_len = (frames // pF) * (height // pH) * (width // pW)
+            item = (
+                item[:ori_len]
+                .view(frames // pF, height // pH, width // pW, pF, pH, pW, self.out_channels)
+                .permute(6, 0, 3, 1, 4, 2, 5)
+                .reshape(self.out_channels, frames, height, width)
+            )
+            if not self.multi_frame_output:
+                item = item[:, :1]
+            out.append(item)
+        return out
+
+    def forward(
+        self,
+        x: List[torch.Tensor],  # per item (C, F, H, W) noisy latents
+        t: torch.Tensor,  # (B,) in the model's convention: 0 = noise, 1 = clean
+        cap_feats: List[torch.Tensor],  # per item (L1, cap_feat_dim) query-token conditions
+        cap_feats_2: Optional[List[torch.Tensor]] = None,  # per item (L2, dim) direct-VLM conditions
+        ref_x: Optional[List[Optional[torch.Tensor]]] = None,  # per item (C, 1, H, W) reference latent
+        patch_size: int = 2,
+        f_patch_size: int = 1,
+        return_dict: bool = True,
+    ):
+        assert patch_size in self.all_patch_size and f_patch_size in self.all_f_patch_size
+        bsz = len(x)
+        device = x[0].device
+        if cap_feats_2 is None:
+            cap_feats_2 = [None] * bsz
+        if ref_x is None:
+            ref_x = [None] * bsz
+
+        adaln_input = self.t_embedder(t * self.t_scale).type_as(x[0])
+        learned_pad = self.alignment_padding_mode == LEARNED_PADDING
+
+        # ---- per-item layout: captions first (positions 1..), then the image.
+        # The 32-alignment pad is materialized only under learned padding; the
+        # zero_masked reference masks those slots out of everything, so here
+        # they exist only as a gap in the position grid.
+        cap_in, cap_direct, cap_pad_lens, cap_pos_ids = [], [], [], []
+        img_feats, img_sizes, img_pos_ids, img_pad_lens = [], [], [], []
+        for image, cap, direct, ref in zip(x, cap_feats, cap_feats_2, ref_x):
+            cap_len = len(cap) + (len(direct) if direct is not None else 0)
+            cap_pad = (-cap_len) % SEQ_MULTI_OF
+            n_pad = cap_pad if learned_pad else 0
+            cap_in.append(cap)
+            cap_direct.append(direct)
+            cap_pad_lens.append(n_pad)
+            cap_pos_ids.append(self._positions(cap_len, 0, (cap_len + n_pad, 1, 1), (1, 0, 0), device))
+
+            if ref is not None:
+                image = torch.cat([image, ref.to(image.dtype)], dim=1)
+            C, Fr, H, W = image.shape
+            F_t, H_t, W_t = Fr // f_patch_size, H // patch_size, W // patch_size
+            patches = (
+                image.view(C, F_t, f_patch_size, H_t, patch_size, W_t, patch_size)
+                .permute(1, 3, 5, 2, 4, 6, 0)
+                .reshape(F_t * H_t * W_t, f_patch_size * patch_size * patch_size * C)
+            )
+            img_len = patches.shape[0]
+            n_pad = (-img_len) % SEQ_MULTI_OF if learned_pad else 0
+            img_feats.append(self._pad_feat(patches, n_pad))
+            img_sizes.append((Fr, H, W))
+            img_pad_lens.append(n_pad)
+            img_pos_ids.append(
+                self._positions(img_len, n_pad, (F_t, H_t, W_t), (cap_len + cap_pad + 1, 0, 0), device)
+            )
+
+        # ---- image tokens: embed, pad, refine
+        x_seqlens = [len(f) for f in img_feats]
+        x_emb = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](torch.cat(img_feats, dim=0))
+        x_emb = self._apply_pad_token(x_emb, img_pad_lens, x_seqlens, self.x_pad_token)
+        x_emb, x_freqs, x_mask, _ = self._batch(list(x_emb.split(x_seqlens, dim=0)), img_pos_ids, device)
+        for layer in self.noise_refiner:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                x_emb = self._gradient_checkpointing_func(layer, x_emb, x_mask, x_freqs, adaln_input)
+            else:
+                x_emb = layer(x_emb, x_mask, x_freqs, adaln_input)
+
+        # ---- caption tokens: the query stream through cap_embedder, the direct
+        # stream appended as-is
+        cap_lens_1 = [len(c) for c in cap_in]
+        cap_emb = list(self.cap_embedder(torch.cat(cap_in, dim=0)).split(cap_lens_1, dim=0))
+        cap_items = []
+        for emb, direct, pad in zip(cap_emb, cap_direct, cap_pad_lens):
+            if direct is not None:
+                emb = torch.cat([emb, direct.to(emb.dtype)], dim=0)
+            cap_items.append(self._pad_feat(emb, pad))
+        cap_seqlens = [len(c) for c in cap_items]
+        cap_cat = self._apply_pad_token(torch.cat(cap_items, dim=0), cap_pad_lens, cap_seqlens, self.cap_pad_token)
+        cap_batch, cap_freqs, cap_mask, _ = self._batch(list(cap_cat.split(cap_seqlens, dim=0)), cap_pos_ids, device)
+        for layer in self.context_refiner:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                cap_batch = self._gradient_checkpointing_func(layer, cap_batch, cap_mask, cap_freqs)
+            else:
+                cap_batch = layer(cap_batch, cap_mask, cap_freqs)
+
+        # ---- unified sequence: [image, caption] per item
+        unified, unified_freqs = [], []
+        for i in range(bsz):
+            x_len, cap_len = x_seqlens[i], cap_seqlens[i]
+            unified.append(torch.cat([x_emb[i][:x_len], cap_batch[i][:cap_len]]))
+            unified_freqs.append(torch.cat([x_freqs[i][:x_len], cap_freqs[i][:cap_len]]))
+        unified_mask = self._length_mask([len(u) for u in unified], device)
+        unified = pad_sequence(unified, batch_first=True, padding_value=0.0)
+        unified_freqs = pad_sequence(unified_freqs, batch_first=True, padding_value=0.0)
+
+        for layer in self.layers:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                unified = self._gradient_checkpointing_func(layer, unified, unified_mask, unified_freqs, adaln_input)
+            else:
+                unified = layer(unified, unified_mask, unified_freqs, adaln_input)
+
+        unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input)
+        out = self.unpatchify(list(unified.unbind(dim=0)), img_sizes, patch_size, f_patch_size)
+        if not return_dict:
+            return (out,)
+        return Transformer2DModelOutput(sample=out)

extensions_built_in/diffusion_models/ming_image/src/vae.py

@@ -0,0 +1,84 @@
+"""The Ming-Image VAE: the Qwen-Image (Wan-style video) VAE with four image
+channels. The ComfyUI repack ships it in Comfy's Wan-native key layout, which
+the load converter maps onto the diffusers module names."""
+
+from toolkit.models.v2.vae.qwen_image import QwenImageVAE
+
+from .checkpoints import BASE_REPO, COMFY_REPO, COMFY_VAE_FILES, comfy_weight_names
+
+# comfy `<block>.residual.<n>` / `shortcut` -> the diffusers resnet submodule
+_RESNET_PARTS = {
+    "residual.0": "norm1",
+    "residual.2": "conv1",
+    "residual.3": "norm2",
+    "residual.6": "conv2",
+    "shortcut": "conv_shortcut",
+}
+# comfy's `middle` Sequential is resnet, attention, resnet
+_MID_PARTS = {"0": "resnets.0", "1": "attentions.0", "2": "resnets.1"}
+
+
+def _resnet_suffix(inner):
+    consumed = 2 if inner[0] == "residual" else 1
+    return [_RESNET_PARTS[".".join(inner[:consumed])]] + inner[consumed:]
+
+
+def comfy_to_diffusers_key(key: str, up_block_len: int) -> str:
+    """One ComfyUI (Wan-native) VAE parameter name -> its diffusers name.
+    `up_block_len` is the number of decoder entries per up block (its resnets
+    plus the upsampler)."""
+    parts = key.split(".")
+    if parts[0] == "conv1":
+        return ".".join(["quant_conv"] + parts[1:])
+    if parts[0] == "conv2":
+        return ".".join(["post_quant_conv"] + parts[1:])
+
+    side, rest = parts[0], parts[1:]
+    if rest[0] == "conv1":
+        return ".".join([side, "conv_in"] + rest[1:])
+    if rest[0] == "head":
+        # head.0 is the output norm, head.2 the output conv (head.1 is SiLU)
+        tail = "norm_out" if rest[1] == "0" else "conv_out"
+        return ".".join([side, tail] + rest[2:])
+    if rest[0] == "middle":
+        part, inner = _MID_PARTS[rest[1]], rest[2:]
+        if part.startswith("resnets"):
+            inner = _resnet_suffix(inner)
+        return ".".join([side, "mid_block", part] + inner)
+
+    index, inner = int(rest[1]), rest[2:]
+    if side == "encoder":
+        # encoder.downsamples.<i>: a flat list of resnets and resamplers in
+        # the same order as diffusers' down_blocks
+        if inner[0] in ("resample", "time_conv"):
+            return ".".join([side, "down_blocks", str(index)] + inner)
+        return ".".join([side, "down_blocks", str(index)] + _resnet_suffix(inner))
+    # decoder.upsamples.<i> is flat too; diffusers nests each up block's
+    # resnets under it with the upsampler last
+    block, j = divmod(index, up_block_len)
+    if inner[0] in ("resample", "time_conv"):
+        return ".".join([side, "up_blocks", str(block), "upsamplers", "0"] + inner)
+    return ".".join([side, "up_blocks", str(block), "resnets", str(j)] + _resnet_suffix(inner))
+
+
+def convert_comfy_vae_state_dict(state_dict):
+    if "encoder.conv1.weight" not in state_dict:
+        return state_dict
+    # the first decoder resampler sits right after a block's resnets
+    first_resample = min(
+        int(k.split(".")[2])
+        for k in state_dict
+        if k.startswith("decoder.upsamples.") and ".resample." in k
+    )
+    up_block_len = first_resample + 1
+    return {comfy_to_diffusers_key(k, up_block_len): v for k, v in state_dict.items()}
+
+
+class MingImageVAE(QwenImageVAE):
+    aitk_config_repo = BASE_REPO
+    aitk_comfy_repo = COMFY_REPO
+    aitk_comfy_weight_names = comfy_weight_names(COMFY_VAE_FILES)
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        return convert_comfy_vae_state_dict(state_dict)

extensions_built_in/diffusion_models/ui.tsx

@@ -620,6 +620,108 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
       },
     ],
   },
+  {
+    name: "ming_image",
+    label: "Ming-Image 0.1 Design (w/ Training Adapter)",
+    group: "image",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      // the ComfyUI repack (local copies in the comfy models folder win over
+      // the download); the vendor repo or a fine-tune here loads as named
+      "config.process[0].model.name_or_path": [
+        "Kijai/Ming-Image-ComfyUI",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      // the text encoder is a 16B MoE (34 GB in bf16): quantize it, and
+      // prefer caching text embeddings so it can be unloaded before training
+      "config.process[0].model.quantize_te": [true, false],
+      // the ComfyUI repack ships int8 convrot for both; these qtypes match it
+      // exactly, so the files attach as-is (another qtype re-quantizes)
+      "config.process[0].model.qtype": ["convrot8", "qfloat8"],
+      "config.process[0].model.qtype_te": ["convrot8", "qfloat8"],
+      "config.process[0].model.low_vram": [true, false],
+      "config.process[0].sample.sampler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.timestep_type": ["shift", "sigmoid"],
+      // training adapter: merged in for training, inverted for sampling
+      "config.process[0].model.assistant_lora_path": [
+        "ostris/ming_image_training_adapter/ming_image_01_design_training_adapter_v1.safetensors",
+        undefined,
+      ],
+      // recommended: 12 steps, no guidance (the negative is zero conditioning)
+      "config.process[0].sample.guidance_scale": [1.0, 4.0],
+      "config.process[0].sample.sample_steps": [12, 25],
+      // the VAE is RGBA: images load, encode and decode with their alpha
+      "config.process[0].model.model_kwargs": [
+        {
+          rgba: false,
+        },
+        {},
+      ],
+    },
+    disableSections: ["network.conv"],
+    // the model also edits (a dataset control path + reference image); the
+    // UI does not expose that yet
+    additionalSections: [
+      "model.low_vram",
+      "model.layer_offloading",
+      "model.assistant_lora_path",
+    ],
+    customModelSelectOptions: [
+      {
+        type: "checkbox",
+        label: "Transparency (RGBA)",
+        getValue: (config: JobConfig) =>
+          config?.config?.process?.[0]?.model?.model_kwargs?.rgba ?? false,
+        onChange: (
+          value: boolean,
+          config: JobConfig,
+          setJobConfig: (value: any, key: string) => void,
+        ) => {
+          const kwargs = {
+            ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}),
+          };
+          if (value) {
+            kwargs.rgba = true;
+          } else {
+            delete kwargs.rgba;
+          }
+          setJobConfig(kwargs, "config.process[0].model.model_kwargs");
+        },
+        doc: {
+          title: "Transparency (RGBA)",
+          description: (
+            <div className="space-y-2">
+              <p>
+                Ming-Image&apos;s VAE is natively RGBA, so alpha can be carried
+                end to end. When on, dataset images and reference images load
+                with their alpha channel, the VAE encodes all four channels, and
+                samples are saved as PNGs with their transparency intact.
+              </p>
+              <p>
+                Images with no alpha of their own get a fully opaque one, so a
+                mixed dataset is fine. Turn this off to train and sample flat
+                RGB: alpha is dropped on the way in and added back as opaque.
+              </p>
+              <p>
+                To sample a transparent image, start the prompt with one of the
+                model&apos;s fixed trigger phrases, e.g.{" "}
+                <code>RGBA, 4-channel, transparent background</code> or{" "}
+                <code>带透明通道，4通道RGBA图像</code>, and sample at 2048x2048:
+                the same prompts come out opaque at 1024.
+              </p>
+              <p>
+                Changing this re-caches latents, and it cannot be combined with
+                a dataset using <code>alpha_mask</code>, which consumes the
+                alpha channel as a loss mask instead.
+              </p>
+            </div>
+          ),
+        },
+      },
+    ],
+  },
   {
     name: "hidream",
     label: "HiDream",

testing/test_model_loading.py

@@ -185,6 +185,20 @@ MODEL_TESTS = {
         "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 1.0},
         "needs_control_image": True,
     },
+    "ming_image": {
+        "model": {
+            "name_or_path": "Kijai/Ming-Image-ComfyUI",
+            "quantize": True,
+            "quantize_te": True,
+            # the repack's int8 convrot files attach as-is at convrot8
+            "qtype": "convrot8",
+            "qtype_te": "convrot8",
+        },
+        # recommended settings: 12 steps, no guidance; the control image
+        # exercises the single-reference editing path
+        "sample": {**IMG, "num_inference_steps": 12, "guidance_scale": 1.0},
+        "needs_control_image": True,
+    },
     # ---- legacy monolith archs (components adopted into v2 on load) ----
     "sd1": {
         "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"},

version.py

@@ -1 +1 @@
-VERSION = "0.13.20"
+VERSION = "0.13.21"