From 77847d7ff83708f0db601bd268e12d3ff4a7fa9e Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 22 Sep 2026 21:25:31 -0600 Subject: [PATCH] Add support for Ming-Image-0.1-Design with training adapter --- README.md | 1 + .../diffusion_models/__init__.py | 2 + .../diffusion_models/ming_image/__init__.py | 3 + .../diffusion_models/ming_image/ming_image.py | 559 ++++++++++++++ .../ming_image/src/__init__.py | 0 .../ming_image/src/bailing_moe_v2.py | 613 +++++++++++++++ .../ming_image/src/checkpoints.py | 34 + .../ming_image/src/moe_kernels.py | 264 +++++++ .../ming_image/src/pipeline.py | 164 ++++ .../ming_image/src/text_encoder.py | 706 ++++++++++++++++++ .../ming_image/src/transformer.py | 591 +++++++++++++++ .../diffusion_models/ming_image/src/vae.py | 84 +++ extensions_built_in/diffusion_models/ui.tsx | 102 +++ testing/test_model_loading.py | 14 + version.py | 2 +- 15 files changed, 3138 insertions(+), 1 deletion(-) create mode 100644 extensions_built_in/diffusion_models/ming_image/__init__.py create mode 100644 extensions_built_in/diffusion_models/ming_image/ming_image.py create mode 100644 extensions_built_in/diffusion_models/ming_image/src/__init__.py create mode 100644 extensions_built_in/diffusion_models/ming_image/src/bailing_moe_v2.py create mode 100644 extensions_built_in/diffusion_models/ming_image/src/checkpoints.py create mode 100644 extensions_built_in/diffusion_models/ming_image/src/moe_kernels.py create mode 100644 extensions_built_in/diffusion_models/ming_image/src/pipeline.py create mode 100644 extensions_built_in/diffusion_models/ming_image/src/text_encoder.py create mode 100644 extensions_built_in/diffusion_models/ming_image/src/transformer.py create mode 100644 extensions_built_in/diffusion_models/ming_image/src/vae.py diff --git a/README.md b/README.md index 3d457de..5fc8c81 100644 --- a/README.md +++ b/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 diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index aeedcf0..660f0e2 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/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, diff --git a/extensions_built_in/diffusion_models/ming_image/__init__.py b/extensions_built_in/diffusion_models/ming_image/__init__.py new file mode 100644 index 0000000..9c4da41 --- /dev/null +++ b/extensions_built_in/diffusion_models/ming_image/__init__.py @@ -0,0 +1,3 @@ +from .ming_image import MingImageModel + +__all__ = ["MingImageModel"] diff --git a/extensions_built_in/diffusion_models/ming_image/ming_image.py b/extensions_built_in/diffusion_models/ming_image/ming_image.py new file mode 100644 index 0000000..63a327d --- /dev/null +++ b/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 diff --git a/extensions_built_in/diffusion_models/ming_image/src/__init__.py b/extensions_built_in/diffusion_models/ming_image/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/diffusion_models/ming_image/src/bailing_moe_v2.py b/extensions_built_in/diffusion_models/ming_image/src/bailing_moe_v2.py new file mode 100644 index 0000000..c897d8c --- /dev/null +++ b/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 + `` 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 diff --git a/extensions_built_in/diffusion_models/ming_image/src/checkpoints.py b/extensions_built_in/diffusion_models/ming_image/src/checkpoints.py new file mode 100644 index 0000000..e317e2c --- /dev/null +++ b/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)} diff --git a/extensions_built_in/diffusion_models/ming_image/src/moe_kernels.py b/extensions_built_in/diffusion_models/ming_image/src/moe_kernels.py new file mode 100644 index 0000000..7cadbdb --- /dev/null +++ b/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 diff --git a/extensions_built_in/diffusion_models/ming_image/src/pipeline.py b/extensions_built_in/diffusion_models/ming_image/src/pipeline.py new file mode 100644 index 0000000..9326383 --- /dev/null +++ b/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) diff --git a/extensions_built_in/diffusion_models/ming_image/src/text_encoder.py b/extensions_built_in/diffusion_models/ming_image/src/text_encoder.py new file mode 100644 index 0000000..82c980c --- /dev/null +++ b/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 + + `` + 256 x `` + `` + -> 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 `...` 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_PATCH = "" +IMAGE_END = "" + +# the vendor processor's template (processing_bailingmm2.py), NOT the +# tokenizer's own chat template: that one drops the assistant persona line +SYSTEM_PROMPT = "SYSTEM你是一个友好的AI助手。\n\ndetailed thinking off" +HUMAN_PREFIX = "HUMAN" +ASSISTANT_PREFIX = "ASSISTANT" + +# 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. + 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 diff --git a/extensions_built_in/diffusion_models/ming_image/src/transformer.py b/extensions_built_in/diffusion_models/ming_image/src/transformer.py new file mode 100644 index 0000000..53a1497 --- /dev/null +++ b/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) diff --git a/extensions_built_in/diffusion_models/ming_image/src/vae.py b/extensions_built_in/diffusion_models/ming_image/src/vae.py new file mode 100644 index 0000000..83558ca --- /dev/null +++ b/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 `.residual.` / `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.: 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. 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) diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index 13f5d5d..d8ef92e 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/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: ( +
+

+ Ming-Image'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. +

+

+ 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. +

+

+ To sample a transparent image, start the prompt with one of the + model's fixed trigger phrases, e.g.{" "} + RGBA, 4-channel, transparent background or{" "} + 带透明通道，4通道RGBA图像, and sample at 2048x2048: + the same prompts come out opaque at 1024. +

+

+ Changing this re-caches latents, and it cannot be combined with + a dataset using alpha_mask, which consumes the + alpha channel as a loss mask instead. +

+
+ ), + }, + }, + ], + }, { name: "hidream", label: "HiDream", diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 108a7c5..f9daf44 100644 --- a/testing/test_model_loading.py +++ b/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"}, diff --git a/version.py b/version.py index d55c7ea..c4f5064 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.13.20" +VERSION = "0.13.21"