diff --git a/.gitignore b/.gitignore index b0456244..558cf816 100644 --- a/.gitignore +++ b/.gitignore @@ -186,4 +186,5 @@ aitk_db.db-wal aitk_db.db-shm /notes.md /data -.claude \ No newline at end of file +.claude +original_repo \ No newline at end of file diff --git a/README.md b/README.md index f0888f33..a27d9376 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,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) +- [microsoft/Mage-Flow-Base](https://huggingface.co/microsoft/Mage-Flow-Base) (Mage-Flow) ### Instruction / Edit - [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) (FLUX.1-Kontext-dev) @@ -44,6 +45,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [Boogu/Boogu-Image-0.1-Edit](https://huggingface.co/Boogu/Boogu-Image-0.1-Edit) (Boogu Image Edit) - [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (Krea 2 Edit Training) - [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (Krea 2 Turbo Edit Training) +- [microsoft/Mage-Flow-Edit-Base](https://huggingface.co/microsoft/Mage-Flow-Edit-Base) (Mage-Flow Edit) ### Video - [Wan-AI/Wan2.1-T2V-1.3B-Diffusers](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers) (Wan 2.1 1.3B) diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index ed10bc70..c7e8a147 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/extensions_built_in/diffusion_models/__init__.py @@ -18,6 +18,7 @@ from .ideogram4 import Ideogram4Model from .prx_pixel_t2i import PRXPixelT2IModel from .krea2 import Krea2Model from .boogu_image import BooguImageModel, BooguImageEditModel +from .mageflow import MageFlowModel, MageFlowEditModel AI_TOOLKIT_MODELS = [ # put a list of models here @@ -51,4 +52,6 @@ AI_TOOLKIT_MODELS = [ Krea2Model, BooguImageModel, BooguImageEditModel, + MageFlowModel, + MageFlowEditModel, ] diff --git a/extensions_built_in/diffusion_models/mageflow/__init__.py b/extensions_built_in/diffusion_models/mageflow/__init__.py new file mode 100644 index 00000000..aca4cbf0 --- /dev/null +++ b/extensions_built_in/diffusion_models/mageflow/__init__.py @@ -0,0 +1 @@ +from .mageflow import MageFlowModel, MageFlowEditModel diff --git a/extensions_built_in/diffusion_models/mageflow/mageflow.py b/extensions_built_in/diffusion_models/mageflow/mageflow.py new file mode 100644 index 00000000..0d15ffd1 --- /dev/null +++ b/extensions_built_in/diffusion_models/mageflow/mageflow.py @@ -0,0 +1,650 @@ +"""Mage-Flow (microsoft/Mage) for ai-toolkit. + +Two archs backed by the same NR-MMDiT stack: + - ``mageflow`` text-to-image (microsoft/Mage-Flow-Base) + - ``mageflow_edit`` instruction edit (microsoft/Mage-Flow-Edit-Base) + +Components (diffusers-style repo layout): + - transformer: ``MageFlow`` dual-stream DiT (src/transformer.py) — packed + variable-length [text | image(+refs)] sequences, per-sample 2D RoPE, + joint varlen attention. 4B params, hidden 3072, 12 blocks. + - text encoder: Qwen3-VL (ships inside the repo under ``text_encoder/``); + conditioning is the final hidden states with the templated system prompt + dropped (34 tokens for t2i, 64 for edit). + - autoencoder: ``MageVAE`` (src/vae.py) — 128-channel, 16x downsample + one-step diffusion codec, no latent normalization. + +Flow matching convention matches ai-toolkit exactly (sigma=1 noise -> sigma=0 +clean, target = noise - clean), static sigma shift 6.0. Edit mode feeds the +reference images in two places, as in the reference implementation: through +the Qwen3-VL encoder alongside the instruction (long edge capped at 384) and +as clean VAE latents sequence-appended after the noisy target tokens (at the +target resolution). +""" + +import json +import math +import os +from typing import TYPE_CHECKING, List, Optional + +import torch +import torch.nn.functional as F +from PIL import Image +from torchvision.transforms.functional import to_tensor +from safetensors.torch import load_file, save_file + +import huggingface_hub +from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration +from optimum.quanto import freeze + +from toolkit.config_modules import GenerateImageConfig, ModelConfig +from toolkit.models.base_model import BaseModel +from toolkit.basic import flush +from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds +from toolkit.samplers.custom_flowmatch_sampler import ( + CustomFlowMatchEulerDiscreteScheduler, +) +from toolkit.accelerator import unwrap_model +from toolkit.metadata import get_meta_for_safetensors +from toolkit.util.quantize import quantize, get_qtype, quantize_model +from toolkit.memory_management import MemoryManager + +from .src.transformer import MageFlow, MageFlowParams +from .src.vae import MageVAE +from .src.text_encoder import ( + encode_mageflow_prompt, + patch_qwen_vl_patch_embed, + resize_vl_images, +) +from .src.pipeline import MageFlowPipeline, predict_velocity + +if TYPE_CHECKING: + from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO + + +# Training/sampling timestep distribution: the released checkpoints use a +# static sigma shift of 6.0 (transformer/config.json + scheduler config). +scheduler_config = { + "num_train_timesteps": 1000, + "use_dynamic_shifting": False, + "shift": 6.0, +} + +# Keys of the checkpoint transformer/config.json that are NOT MageFlowParams +# constructor args (legacy/unused fields) — same filter as the reference +# ``load_from_repo``. Everything else becomes the DiT structure. +_CONFIG_META_KEYS = { + "_class_name", + "txt_max_length", + "max_sequence_length", + "param_dtype", + "packing", + "schedule_mode", + "static_shift", + "use_time_shift", + "rope_type", + "apply_text_rotary_emb", + "mlp_ratio", + "depth_single_blocks", + "theta", + "qkv_bias", + "guidance_embed", + "vec_in_dim", + "vec_type", + "time_type", + "double_block_type", +} + +HF_TOKEN = os.getenv("HF_TOKEN", None) + + +class MageFlowModel(BaseModel): + arch = "mageflow" + is_edit = False + use_old_lokr_format = False + + 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.target_lora_modules = ["MageFlow"] + + # MageVAE is 16x downsampling; the DiT patch size is 1 (one token per + # latent pixel), so pixel sizes must be multiples of 16. + self.patch_size = 1 + self.vae_scale_factor = 16 + # Safety cap on prompt token length (truncation only); embeds are stored + # per-sample at natural length and packed varlen at the model call. + self.max_text_length = int( + self.model_config.model_kwargs.get("max_text_length", 2048) + ) + # Qwen3-VL AutoProcessor for encoding reference images into the prompt + # (edit only). + self.vl_processor = None + + if self.is_edit: + # Reference images feed the model in two places: through the + # Qwen3-VL encoder alongside the instruction, and as clean VAE + # latents sequence-appended after the noisy target tokens. + self.encode_control_in_text_embeddings = True + self.has_multiple_control_images = True + # References keep their own aspect/size in the dataloader; they are + # resized to the target resolution here (reference behavior). + self.use_raw_control_images = True + + @staticmethod + def get_train_scheduler(): + return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config) + + def get_bucket_divisibility(self): + # 16 for the VAE downsample, patch size is 1. + return self.vae_scale_factor * self.patch_size + + # ------------------------------------------------------------------ + # Loading + # ------------------------------------------------------------------ + def _get_model_file(self, relpath: str) -> str: + """Resolve a repo-relative file from a local directory or the HF hub.""" + name_or_path = self.model_config.name_or_path + if os.path.isdir(name_or_path): + path = os.path.join(name_or_path, relpath) + if not os.path.exists(path): + raise FileNotFoundError( + f"Could not find {relpath!r} in {name_or_path!r}" + ) + return path + return huggingface_hub.hf_hub_download( + repo_id=name_or_path, filename=relpath, token=HF_TOKEN + ) + + def _load_transformer(self) -> MageFlow: + dtype = self.torch_dtype + self.print_and_status_update("Loading transformer (MageFlow NR-MMDiT)") + + with open(self._get_model_file("transformer/config.json")) as f: + tcfg = json.load(f) + structure = {k: v for k, v in tcfg.items() if k not in _CONFIG_META_KEYS} + structure.update(self.model_config.model_kwargs.get("transformer_config", {})) + params = MageFlowParams(**structure) + + # Build on meta, then materialize straight from the checkpoint. + with torch.device("meta"): + transformer = MageFlow(params) + + self.print_and_status_update(" - fetching transformer weights") + state_dict = load_file( + self._get_model_file("transformer/diffusion_pytorch_model.safetensors") + ) + state_dict = { + k: (v.to(dtype) if v.is_floating_point() else v) + for k, v in state_dict.items() + } + self.print_and_status_update(" - loading transformer state dict") + transformer.load_state_dict(state_dict, strict=True, assign=True) + # The RoPE tables are plain tensors (not buffers/params), so the meta + # init left them unmaterialized — rebuild them for real. + transformer.reset_rope() + del state_dict + flush() + return transformer + + def _load_text_encoder(self): + dtype = self.torch_dtype + te_path = self.model_config.model_kwargs.get("text_encoder_path", None) + if te_path is not None: + te_kwargs = {} + else: + te_path = self.model_config.name_or_path + te_kwargs = {"subfolder": "text_encoder"} + self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}") + + tokenizer = AutoTokenizer.from_pretrained(te_path, token=HF_TOKEN, **te_kwargs) + text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + te_path, torch_dtype=dtype, token=HF_TOKEN, **te_kwargs + ) + vl_processor = None + if self.is_edit: + # Edit mode: reference images are encoded into the text embeddings, + # so the vision tower stays. Swap its Conv3d patch_embed for an + # equivalent GEMM (bf16 Conv3d has no fast cuDNN kernel). + vl_processor = AutoProcessor.from_pretrained( + te_path, token=HF_TOKEN, **te_kwargs + ) + patch_qwen_vl_patch_embed(text_encoder) + else: + # We only ever encode text, so the vision tower is dead weight -- + # drop it to free VRAM. + if getattr(text_encoder.model, "visual", None) is not None: + text_encoder.model.visual = None + text_encoder.eval() + text_encoder.requires_grad_(False) + flush() + return tokenizer, vl_processor, text_encoder + + def _load_vae(self) -> MageVAE: + self.print_and_status_update("Loading MageVAE") + vae_path = self.model_config.model_kwargs.get("vae_path", None) + if vae_path is None: + vae_path = self._get_model_file("vae/diffusion_pytorch_model.safetensors") + vae = MageVAE( + ckpt_path=vae_path, + sample_posterior=bool( + self.model_config.model_kwargs.get("vae_sample_posterior", True) + ), + ) + vae.eval() + vae.requires_grad_(False) + return vae + + def get_quantization_exclude_modules(self): + # sensitive modules kept in full precision (fnmatch patterns on module + # names within MageFlow): + # img_in / txt_in / txt_norm - input projections + # time_text_embed* - timestep embedder feeding every + # block's modulation + # norm_out* / proj_out - final adaptive norm / output projection + return [ + "img_in", + "txt_in", + "txt_norm", + "time_text_embed*", + "norm_out*", + "proj_out", + ] + + def load_model(self): + dtype = self.torch_dtype + self.print_and_status_update("Loading Mage-Flow model") + + transformer = self._load_transformer() + + if self.model_config.quantize: + self.print_and_status_update("Quantizing transformer") + quantize_model(self, transformer) + flush() + + if ( + self.model_config.layer_offloading + and self.model_config.layer_offloading_transformer_percent > 0 + ): + MemoryManager.attach( + transformer, + self.device_torch, + offload_percent=self.model_config.layer_offloading_transformer_percent, + ) + + if self.model_config.low_vram: + self.print_and_status_update("Moving transformer to CPU") + transformer.to("cpu") + else: + transformer.to(self.device_torch, dtype=dtype) + flush() + + tokenizer, vl_processor, text_encoder = self._load_text_encoder() + if self.model_config.quantize_te: + self.print_and_status_update("Quantizing text encoder") + text_encoder.to(self.device_torch) + quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) + freeze(text_encoder) + flush() + if ( + self.model_config.layer_offloading + and self.model_config.layer_offloading_text_encoder_percent > 0 + ): + MemoryManager.attach( + text_encoder, + self.device_torch, + offload_percent=self.model_config.layer_offloading_text_encoder_percent, + ) + + if self.model_config.low_vram: + self.print_and_status_update("Moving text encoder to CPU") + text_encoder.to("cpu") + else: + text_encoder.to(self.device_torch) + flush() + + vae = self._load_vae() + vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype) + + self.noise_scheduler = MageFlowModel.get_train_scheduler() + + self.vae = vae + self.text_encoder = text_encoder + self.tokenizer = tokenizer + self.vl_processor = vl_processor + self.model = transformer + self.pipeline = MageFlowPipeline(self) + self.print_and_status_update("Model Loaded") + + # ------------------------------------------------------------------ + # Generation (training previews) + # ------------------------------------------------------------------ + def get_generation_pipeline(self): + return MageFlowPipeline(self) + + def generate_single_image( + self, + pipeline: MageFlowPipeline, + 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) + + sc = self.get_bucket_divisibility() + gen_config.width = int(gen_config.width // sc * sc) + gen_config.height = int(gen_config.height // sc * sc) + + # Reference image(s) -> clean VAE latents appended to the sequence. + # The Qwen3-VL side already saw them (baked into the prompt embeds). + # ctrl_img_1 mirrors ctrl_img when unset, so use one or the other. + ctrl_paths = [] + if self.is_edit: + if gen_config.ctrl_img is not None: + ctrl_paths.append(gen_config.ctrl_img) + elif gen_config.ctrl_img_1 is not None: + ctrl_paths.append(gen_config.ctrl_img_1) + if gen_config.ctrl_img_2 is not None: + ctrl_paths.append(gen_config.ctrl_img_2) + if gen_config.ctrl_img_3 is not None: + ctrl_paths.append(gen_config.ctrl_img_3) + + ref_latents = None + if ctrl_paths: + ctrl_tensors = [ + to_tensor(Image.open(path).convert("RGB")) for path in ctrl_paths + ] + # one batch item (preview batch size is 1) -> List[List[(128, h, w)]] + ref_latents = [ + self._encode_ref_latents( + ctrl_tensors, gen_config.height, gen_config.width + ) + ] + + img = 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, + ref_latents=ref_latents, + )[0] + return img + + # ------------------------------------------------------------------ + # Reference-image helpers (edit) + # ------------------------------------------------------------------ + def _encode_ref_latents( + self, control_tensors, target_height: int, target_width: int + ) -> List[torch.Tensor]: + """Encode ``[0, 1]`` reference image tensors to clean VAE latents. + + Sizing (deliberate deviation from the reference ``generate_edits``, + which resizes every ref to exactly the target resolution and squishes + mismatched aspect ratios): each reference is resized to the TARGET's + total pixel count while keeping its OWN aspect ratio, snapped to the + 16px divisibility. A ref that shares the target's aspect ratio still + lands on exactly the target size (identical to the reference + behavior); one with a different aspect ratio keeps its shape instead + of being distorted. Returns a list of ``(128, h, w)`` latents (one per + reference image). ``control_tensors`` is a list of ``(C, H, W)`` or + ``(1, C, H, W)`` tensors in ``[0, 1]``. + """ + sc = self.get_bucket_divisibility() # 16 + target_area = target_height * target_width + latents = [] + for img in control_tensors: + if img.dim() == 3: + img = img.unsqueeze(0) + img = img.to(self.device_torch, dtype=torch.float32) + + h, w = img.shape[2], img.shape[3] + ratio = h / w + new_h = math.sqrt(target_area * ratio) + new_w = new_h / ratio + new_h = max(sc, int(round(new_h / sc)) * sc) + new_w = max(sc, int(round(new_w / sc)) * sc) + if (new_h, new_w) != (h, w): + img = F.interpolate( + img, + size=(new_h, new_w), + mode="bicubic", + antialias=True, + ).clamp(0, 1) + # encode_images expects [-1, 1]; control tensors arrive in [0, 1]. + latent = self.encode_images( + img * 2 - 1, device=self.device_torch, dtype=self.torch_dtype + ) + latents.append(latent[0]) # drop batch dim -> (128, h, w) + return latents + + def _batch_ref_latents_from_batch( + self, + batch: "DataLoaderBatchDTO", + batch_size: int, + target_height: int, + target_width: int, + ) -> Optional[List[List[torch.Tensor]]]: + """Build predict_velocity's ``ref_latents`` from a train batch.""" + control_list = batch.control_tensor_list + if control_list is None and batch.control_tensor is not None: + control_list = [batch.control_tensor[b : b + 1] for b in range(batch_size)] + if control_list is None: + return None + if len(control_list) != batch_size: + raise ValueError("Control tensor list length does not match batch size") + ref_latents = [] + for controls in control_list: + if isinstance(controls, torch.Tensor): + controls = [controls] + ref_latents.append( + self._encode_ref_latents(controls, target_height, target_width) + ) + return ref_latents + + # ------------------------------------------------------------------ + # Training hooks + # ------------------------------------------------------------------ + def get_noise_prediction( + self, + latent_model_input: torch.Tensor, # (B, 128, 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) + + # Clean reference latents from the batch's control images (if any); + # they ride along in the sequence, never noised. + ref_latents = None + if batch is not None and self.is_edit: + with torch.no_grad(): + _, _, lh, lw = latent_model_input.shape + ref_latents = self._batch_ref_latents_from_batch( + batch, + latent_model_input.shape[0], + target_height=lh * self.vae_scale_factor, + target_width=lw * self.vae_scale_factor, + ) + + # toolkit timestep (0..1000, 1000 = pure noise) -> flow sigma in + # [0, 1] with 1 = pure noise. Same convention -> straight divide. + t = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0 + if t.dim() == 0: + t = t.unsqueeze(0) + if t.shape[0] != latent_model_input.shape[0]: + t = t.expand(latent_model_input.shape[0]) + + pred = predict_velocity( + self.transformer, + latent_model_input.to(self.device_torch, self.torch_dtype), + t, + text_embeddings.text_embeds, + ref_latents=ref_latents, + ) + return pred + + def get_prompt_embeds(self, prompt, control_images=None) -> AdvancedPromptEmbeds: + if isinstance(prompt, str): + prompt = [prompt] + + if self.text_encoder.device == torch.device("cpu"): + self.text_encoder.to(self.device_torch) + + # Normalize control images to a per-prompt list (List[List[Tensor]]). + # They arrive as a (B, C, H, W) batch tensor (control_tensor), a list of + # per-sample lists (control_tensor_list), or a flat list of (1, C, H, W) + # tensors for a single prompt (sampling / blank-embed caching). + if control_images is not None: + if isinstance(control_images, torch.Tensor): + control_images = [ + [control_images[i]] for i in range(control_images.shape[0]) + ] + elif len(control_images) > 0 and not isinstance(control_images[0], list): + control_images = [control_images] + if len(control_images) == 1 and len(prompt) > 1: + control_images = control_images * len(prompt) + if len(control_images) != len(prompt): + raise ValueError( + "Number of prompts must match number of control image sets" + ) + else: + control_images = [None] * len(prompt) + + template_name = "mage-flow-edit" if self.is_edit else "mage-flow" + vl_long_edge = int(self.model_config.model_kwargs.get("vl_cond_long_edge", 384)) + + # Encode each prompt at its natural length and store one (L, 2560) + # tensor per batch item. Padding is never needed: the packed varlen + # model call consumes the per-sample lengths directly. + features_list = [] + for p, ctrl in zip(prompt, control_images): + images = None + if self.is_edit and ctrl is not None and len(ctrl) > 0: + images = resize_vl_images(ctrl, max_long_edge=vl_long_edge) + features = encode_mageflow_prompt( + self.text_encoder, + self.tokenizer, + p, + template_name=template_name, + max_length=self.max_text_length, + images=images, + processor=self.vl_processor, + dtype=self.torch_dtype, + ) + features_list.append(features) + + return AdvancedPromptEmbeds(text_embeds=features_list) + + 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 + + # ------------------------------------------------------------------ + # VAE (MageVAE -- raw latents, no normalization) + # ------------------------------------------------------------------ + def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None): + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + self.vae.eval() + + if isinstance(image_list, list): + images = torch.stack( + [img.squeeze(0) if img.dim() == 4 else img for img in image_list], dim=0 + ) + else: + images = image_list + images = images.to(self.vae.device, dtype=self.vae.dtype) + + latents = self.vae.encode(images) # (B, 128, H/16, W/16) + return latents.to(device, dtype=dtype) + + def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + + latents = latents.to(self.vae.device, dtype=self.vae.dtype) + images = self.vae.decode(latents) # (B, 3, H, W) in [-1, 1] + return images.to(device, dtype=dtype) + + # ------------------------------------------------------------------ + # Saving / bookkeeping + # ------------------------------------------------------------------ + def save_model(self, output_path, meta, save_dtype): + from toolkit.util.quantize import dequantize_if_quantized + + if not output_path.endswith(".safetensors"): + output_path = output_path + ".safetensors" + transformer: MageFlow = unwrap_model(self.model) + state_dict = transformer.state_dict() + save_dict = {} + for k, v in state_dict.items(): + # dequantize any quantized (e.g. quanto/torchao) weights so we save plain full precision tensors + save_dict[k] = ( + dequantize_if_quantized(v).clone().to("cpu", dtype=save_dtype) + ) + meta = get_meta_for_safetensors(meta, name=self.arch) + save_file(save_dict, output_path, metadata=meta) + + def get_base_model_version(self): + return "mageflow" + + def get_transformer_block_names(self) -> Optional[List[str]]: + return ["transformer_blocks"] + + def convert_lora_weights_before_save(self, state_dict): + return { + k.replace("transformer.", "diffusion_model."): v + for k, v in state_dict.items() + } + + def convert_lora_weights_before_load(self, state_dict): + return { + k.replace("diffusion_model.", "transformer."): v + for k, v in state_dict.items() + } + + +class MageFlowEditModel(MageFlowModel): + arch = "mageflow_edit" + is_edit = True + + def get_base_model_version(self): + return "mageflow_edit" diff --git a/extensions_built_in/diffusion_models/mageflow/src/__init__.py b/extensions_built_in/diffusion_models/mageflow/src/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/extensions_built_in/diffusion_models/mageflow/src/attn.py b/extensions_built_in/diffusion_models/mageflow/src/attn.py new file mode 100644 index 00000000..d768602a --- /dev/null +++ b/extensions_built_in/diffusion_models/mageflow/src/attn.py @@ -0,0 +1,73 @@ +"""Varlen attention shim for Mage-Flow. + +The reference implementation (github.com/microsoft/Mage) runs every attention +call through flash-attn's ``flash_attn_varlen_func`` so that several +variable-length samples can be packed into one sequence and kept isolated via +``cu_seqlens``. This shim exposes the same function with the FA2 calling +convention, falling back to a per-sequence ``torch.scaled_dot_product_attention`` +loop when flash-attn is not installed (functionally equivalent for the dense / +non-causal path the DiT uses, just slower). +""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch +import torch.nn.functional as F + +_RESOLVED_FN: Callable[..., Any] | None = None + + +def _sdpa_varlen( + q, + k, + v, + cu_seqlens_q=None, + cu_seqlens_k=None, + max_seqlen_q=None, + max_seqlen_k=None, + dropout_p: float = 0.0, + softmax_scale=None, + causal: bool = False, + **_unused: Any, +): + """FA2-varlen-compatible SDPA fallback: one SDPA dispatch per sequence.""" + cu_q = cu_seqlens_q.tolist() + cu_k = cu_seqlens_k.tolist() + outs = [] + for qs, qe, ks, ke in zip(cu_q[:-1], cu_q[1:], cu_k[:-1], cu_k[1:]): + # (s, h, d) -> (1, h, s, d) + q_i = q[qs:qe].transpose(0, 1).unsqueeze(0) + k_i = k[ks:ke].transpose(0, 1).unsqueeze(0) + v_i = v[ks:ke].transpose(0, 1).unsqueeze(0) + out_i = F.scaled_dot_product_attention( + q_i, + k_i, + v_i, + attn_mask=None, + dropout_p=dropout_p, + is_causal=causal, + scale=softmax_scale, + ) + outs.append(out_i.squeeze(0).transpose(0, 1)) + return torch.cat(outs, dim=0).contiguous() + + +def _resolve() -> Callable[..., Any]: + global _RESOLVED_FN + if _RESOLVED_FN is None: + try: + from flash_attn import flash_attn_varlen_func as _fn + + _RESOLVED_FN = _fn + except ImportError: + _RESOLVED_FN = _sdpa_varlen + return _RESOLVED_FN + + +def flash_attn_varlen_func(*args, **kwargs): + return _resolve()(*args, **kwargs) + + +__all__ = ["flash_attn_varlen_func"] diff --git a/extensions_built_in/diffusion_models/mageflow/src/pipeline.py b/extensions_built_in/diffusion_models/mageflow/src/pipeline.py new file mode 100644 index 00000000..56159e3c --- /dev/null +++ b/extensions_built_in/diffusion_models/mageflow/src/pipeline.py @@ -0,0 +1,229 @@ +"""Packing / sampling helpers for Mage-Flow. + +Turns latents + Qwen3-VL text features into the packed variable-length +sequences the ``MageFlow`` DiT consumes (batch dim 1, per-sample +``cu_seqlens``, exactly mirroring the reference training/inference packing), +and provides a minimal flow-matching sampler for ai-toolkit preview images. + +Time convention: rectified flow with ``x_t = (1 - sigma) * clean + sigma * +noise``, the model predicts the velocity ``noise - clean``, and sampling +integrates sigma from 1 (pure noise) down to 0 (clean). This is identical to +ai-toolkit's convention, so the toolkit ``timestep / 1000`` flows straight +through as sigma. The inference sigma schedule is the reference's static +shift: ``shifted = shift * s / (1 + (shift - 1) * s)`` over ``s = +linspace(1, 1/steps, steps)`` with a terminal 0 (shift 6.0). + +Edit conditioning: each sample's sequence is ``[target, ref_1, …, ref_N]`` — +clean reference latents ride along after the noisy target tokens, sharing the +sample's timestep modulation. RoPE gives every image segment in the pack its +own "frame" index (segment order), so references land on later frame +coordinates than their target, as in the reference ``generate_edits``. +""" + +from __future__ import annotations + +from typing import List, Optional + +import torch +from einops import rearrange +from PIL import Image +from diffusers.utils.torch_utils import randn_tensor + +from .transformer import MageFlow + + +def lens_to_cu(lens: List[int], device) -> torch.Tensor: + """Sequence lengths -> cumulative cu_seqlens [0, l0, l0+l1, ...] (int32).""" + t = torch.tensor(lens, device=device, dtype=torch.int32) + return torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=device), + torch.cumsum(t, dim=0, dtype=torch.int32), + ] + ) + + +def pack_text_features( + features_list: List[torch.Tensor], device, dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Concatenate per-sample ``(L_i, D)`` text features into the packed + ``[1, sum L, D]`` text stream + its ``cu_seqlens``. No padding — varlen + attention isolates the samples.""" + feats = [f.to(device, dtype) for f in features_list] + txt = torch.cat(feats, dim=0).unsqueeze(0) + txt_cu = lens_to_cu([f.shape[0] for f in feats], device) + return txt, txt_cu + + +def predict_velocity( + model: MageFlow, + latents: torch.Tensor, # (B, C, h, w) noisy target latents + t: torch.Tensor, # (B,) flow sigma in [0, 1] (1 = pure noise) + text_embeds: List[torch.Tensor], # per-sample (L_i, D) Qwen3-VL features + ref_latents: Optional[ + List[List[torch.Tensor]] + ] = None, # per-sample clean (C, hr, wr) refs +) -> torch.Tensor: + """Run the DiT on the packed [text | target(+refs)] sequences. + + Packs the batch into one varlen forward (mirroring reference packing: + per-sample cu_seqlens isolate the samples inside the attention kernel). + Reference latents, when given, are appended clean after each sample's + target tokens; the returned velocity covers only the target tokens, + reshaped back to ``(B, C, h, w)``. No time flip / negation: Mage-Flow's + convention matches ai-toolkit's. + """ + device = latents.device + dtype = latents.dtype + b, c, h, w = latents.shape + + if ref_latents is not None and not any(len(r) > 0 for r in ref_latents): + ref_latents = None + + img_parts, samp_lens, shape_seq, target_idx_parts = [], [], [], [] + off = 0 + for i in range(b): + tgt = rearrange(latents[i], "c h w -> (h w) c") + img_parts.append(tgt) + shape_seq.append((1, h, w)) + target_idx_parts.append(torch.arange(off, off + h * w, device=device)) + samp_len = h * w + off += h * w + if ref_latents is not None: + for ref in ref_latents[i]: + ref = ref.to(device, dtype) + _, rh, rw = ref.shape + img_parts.append(rearrange(ref, "c h w -> (h w) c")) + shape_seq.append((1, rh, rw)) + samp_len += rh * rw + off += rh * rw + samp_lens.append(samp_len) + + img = torch.cat(img_parts, dim=0).unsqueeze(0) # [1, sum, C] + img_cu = lens_to_cu(samp_lens, device) + img_shapes = [shape_seq] + target_idx = torch.cat(target_idx_parts) + + txt, txt_cu = pack_text_features(text_embeds, device, dtype) + + out = model( + img=img, + txt=txt, + timesteps=t.to(device), + img_shapes=img_shapes, + img_cu_seqlens=img_cu, + txt_cu_seqlens=txt_cu, + ) # [1, sum, C_out] + + pred = out[:, target_idx, :] # [1, B*h*w, C] + return rearrange(pred, "1 (b h w) c -> b c h w", b=b, h=h, w=w) + + +def build_shifted_sigmas( + num_steps: int, shift: float = 6.0, device=None +) -> torch.Tensor: + """The reference inference sigma schedule: ``linspace(1, 1/steps, steps)`` + run through the static shift ``shift*s/(1+(shift-1)*s)`` with a terminal 0.""" + s = torch.linspace(1.0, 1.0 / num_steps, num_steps, dtype=torch.float64) + s = shift * s / (1 + (shift - 1) * s) + sigmas = torch.cat([s, torch.zeros(1, dtype=torch.float64)]) + return sigmas.to(device=device, dtype=torch.float32) + + +class MageFlowPipeline: + """Lightweight flow-matching sampler used by ai-toolkit's preview generation.""" + + def __init__(self, model): + # ``model`` is the MageFlowModel (BaseModel) so we can reuse its + # encode/decode and config. + 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, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 30, + guidance_scale: float = 5.0, + latents: Optional[torch.Tensor] = None, + generator: Optional[torch.Generator] = None, + ref_latents: Optional[List[List[torch.Tensor]]] = None, + **kwargs, + ) -> List[Image.Image]: + model = self.model + device = model.device_torch + dtype = model.torch_dtype + transformer: MageFlow = model.transformer + + shift = float(model.model_config.model_kwargs.get("static_shift", 6.0)) + renorm = bool(model.model_config.model_kwargs.get("cfg_renormalization", False)) + + do_cfg = guidance_scale > 1.0 and unconditional_embeds is not None + + gh = height // model.vae_scale_factor + gw = width // model.vae_scale_factor + latent_channels = transformer.in_channels + + if latents is None: + latents = randn_tensor( + (1, latent_channels, gh, gw), + generator=generator, + device=device, + dtype=torch.float32, + ) + latents = latents.to(device, dtype=torch.float32) + + cond_feats = [f.to(device, dtype) for f in conditional_embeds.text_embeds] + if do_cfg: + uncond_feats = [ + f.to(device, dtype) for f in unconditional_embeds.text_embeds + ] + + sigmas = build_shifted_sigmas(num_inference_steps, shift=shift, device=device) + + # Euler integration of the flow ODE, sigma 1 -> 0 (with optional CFG: + # v = uncond + cfg * (cond - uncond), reference convention). + for i in range(num_inference_steps): + s_cur = sigmas[i].item() + s_next = sigmas[i + 1].item() + t = torch.full((latents.shape[0],), s_cur, dtype=dtype, device=device) + v_cond = predict_velocity( + transformer, latents.to(dtype), t, cond_feats, ref_latents=ref_latents + ) + if do_cfg: + v_uncond = predict_velocity( + transformer, + latents.to(dtype), + t, + uncond_feats, + ref_latents=ref_latents, + ) + v = v_uncond + guidance_scale * (v_cond - v_uncond) + if renorm: + # Rescale the guided velocity per token back to the + # conditional velocity's norm (reduces oversaturation). + v = v * ( + torch.norm(v_cond, dim=1, keepdim=True) + / (torch.norm(v, dim=1, keepdim=True) + 1e-6) + ) + else: + v = v_cond + latents = latents + (s_next - s_cur) * v.to(torch.float32) + + images = model.decode_latents(latents, device=device, dtype=dtype) + images = images.float().clamp(-1.0, 1.0) + images = ((images + 1.0) * 127.5).round().to(torch.uint8) + images = images.permute(0, 2, 3, 1).cpu().numpy() + return [Image.fromarray(arr) for arr in images] diff --git a/extensions_built_in/diffusion_models/mageflow/src/text_encoder.py b/extensions_built_in/diffusion_models/mageflow/src/text_encoder.py new file mode 100644 index 00000000..4f3508a2 --- /dev/null +++ b/extensions_built_in/diffusion_models/mageflow/src/text_encoder.py @@ -0,0 +1,169 @@ +"""Qwen3-VL prompt encoding for Mage-Flow. + +The reference implementation (mage_flow/models/modules/text_encoder.py) packs +many prompts into one varlen forward; per-sequence the math reduces to a plain +causal forward with positions ``arange(L)`` replicated on every mrope axis and +the leading system-prompt tokens dropped from the output. ai-toolkit encodes +prompts one at a time (embeds are cached per caption), so that per-sequence +form is implemented here directly on the stock HF Qwen3-VL module. + +The DiT consumes the final hidden states ``[L - drop_idx, 2560]``; the pooled +"vec" of the reference is unused by the released checkpoints (vec_type null — +the transformer adds a zero vector), so it is not computed. +""" + +from __future__ import annotations + +import math +from typing import List, Optional + +import torch +import torch.nn.functional as F + + +# Prompt templates from the reference (mage_flow/models/utils.py). ``start_idx`` +# is the number of leading (system prompt) tokens dropped from the encoder +# output before it conditions the DiT. +PROMPT_TEMPLATE = { + "mage-flow": { + "template": ( + "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, " + "text, spatial relationships of the objects and background:" + "<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + ), + "start_idx": 34, + }, + "mage-flow-edit": { + "template": ( + "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture," + " objects, background), then explain how the user's text instruction should alter or modify the image. " + "Generate a new image that meets the user's requirements while maintaining consistency with the original " + "input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + ), + "start_idx": 64, + }, +} + +# Fixed image placeholder used at edit training time (one per reference image). +EDIT_IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>" + + +def edit_prompt_body(instruction: str, num_refs: int) -> str: + """Training-time multi-reference prompt body: + ``Image 1: Image 2: …{instruction}``.""" + prefix = "".join( + f"Image {j}: {EDIT_IMAGE_PLACEHOLDER}" for j in range(1, num_refs + 1) + ) + return prefix + instruction + + +def patch_qwen_vl_patch_embed(model): + """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain + linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and + falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear + (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves. + Returns the number of patch_embed modules patched. (Same patch as the krea2 + extension / Qwen3VLCaptioner.)""" + patched = 0 + for module in model.modules(): + proj = getattr(module, "proj", None) + if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple( + proj.stride + ): + + def fast_forward(hidden_states, _proj=proj): + w = _proj.weight.reshape(_proj.weight.shape[0], -1) + x = hidden_states.view(-1, w.shape[1]).to(w.dtype) + return F.linear(x, w, _proj.bias) + + module.forward = fast_forward + patched += 1 + return patched + + +def resize_vl_images( + images: List[torch.Tensor], max_long_edge: int = 384 +) -> List["PIL.Image.Image"]: + """Prepare reference images for the Qwen3-VL conditioning pass. + + Matches the reference's ``_resize_long_edge`` (cap the long edge at + ``max_long_edge``, preserving aspect ratio, BICUBIC, never upscale) — the + MLLM only needs a coarse view of the reference; full-resolution detail + flows through the VAE reference latents. Input tensors are ``(C, H, W)`` + or ``(1, C, H, W)`` in [0, 1]; output is PIL for the Qwen processor. + """ + from torchvision.transforms.functional import to_pil_image + + out = [] + for img in images: + if img.dim() == 4: + img = img[0] + img = img.float().clamp(0, 1).cpu() + pil = to_pil_image(img) + if max_long_edge and max_long_edge > 0: + w, h = pil.size + long_edge = max(w, h) + if long_edge > max_long_edge: + scale = max_long_edge / long_edge + new_w = max(1, int(round(w * scale))) + new_h = max(1, int(round(h * scale))) + pil = pil.resize((new_w, new_h), resample=3) # PIL.Image.BICUBIC + out.append(pil) + return out + + +@torch.no_grad() +def encode_mageflow_prompt( + text_encoder, # Qwen3VLForConditionalGeneration + tokenizer, + prompt: str, + template_name: str = "mage-flow", + max_length: int = 2048, + images: Optional[list] = None, # list of PIL images (already VL-resized) + processor=None, # AutoProcessor, required when images are given + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Encode one prompt (optionally with reference images) to DiT conditioning. + + Returns ``(L - start_idx, 2560)`` final hidden states with the system + prompt dropped, matching the reference ``TextEncoder.forward`` txt output. + """ + info = PROMPT_TEMPLATE[template_name] + template = info["template"] + drop_idx = int(info["start_idx"]) + device = next(text_encoder.parameters()).device + + if images: + if processor is None: + raise ValueError( + "encoding reference images requires the Qwen3-VL AutoProcessor" + ) + formatted = template.format(edit_prompt_body(prompt, len(images))) + vl = processor(text=[formatted], images=list(images), return_tensors="pt") + input_ids = vl["input_ids"].to(device) + extra = {} + for key in ("pixel_values", "image_grid_thw"): + if vl.get(key) is not None: + extra[key] = vl[key].to(device) + else: + formatted = template.format(prompt) + input_ids = tokenizer( + formatted, + max_length=max_length + drop_idx, + truncation=True, + return_tensors="pt", + ).input_ids.to(device) + extra = {} + + # The reference packed encoder gives every sequence positions arange(L) + # replicated on all mrope axes (even for image tokens); a 2D position_ids + # is expanded to exactly that by the HF text model. + position_ids = torch.arange(input_ids.shape[1], device=device).unsqueeze(0) + + out = text_encoder.model( + input_ids=input_ids, + position_ids=position_ids, + **extra, + ) + hidden = out.last_hidden_state[0] # (L, D) + return hidden[drop_idx:].to(dtype) diff --git a/extensions_built_in/diffusion_models/mageflow/src/transformer.py b/extensions_built_in/diffusion_models/mageflow/src/transformer.py new file mode 100644 index 00000000..6b0d247c --- /dev/null +++ b/extensions_built_in/diffusion_models/mageflow/src/transformer.py @@ -0,0 +1,809 @@ +"""Mage-Flow NR-MMDiT (dual-stream, native-resolution packed) for ai-toolkit. + +Vendored from the reference implementation +(github.com/microsoft/Mage, mage_flow/models/mage_flow.py + modules/mage_layers.py) +with these deviations: + + - loguru / pydantic / torch._dynamo decorators dropped, + - gradient checkpointing is gated on ``self.gradient_checkpointing and + torch.is_grad_enabled()`` (instead of the reference's ``self.training and + self.checkpoint``) and exposed via ``enable_gradient_checkpointing`` / + ``disable_gradient_checkpointing`` so ai-toolkit's trainer can toggle it, + - flash-attn is imported through the local shim (``.attn``) which falls back + to a per-sequence SDPA loop when flash-attn is unavailable. + +State-dict keys are identical to the reference checkpoints +(``transformer/diffusion_pytorch_model.safetensors`` in +microsoft/Mage-Flow-Base and microsoft/Mage-Flow-Edit-Base). + +Sequence layout: images and text are packed as variable-length sequences of +batch dim 1 (``[1, sum_len, C]``) with per-sample ``cu_seqlens``; joint text+image +attention runs in a single varlen kernel. 2D multi-scale RoPE is applied to +image tokens only (text tokens are not rotated). ``timesteps`` is the flow +sigma in [0, 1] (1 = pure noise), one entry per packed sample. +""" + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +from torch import Tensor +from torch.utils.checkpoint import checkpoint + +from diffusers.models.attention import FeedForward +from diffusers.models.embeddings import TimestepEmbedding +from diffusers.models.normalization import RMSNorm + +from .attn import flash_attn_varlen_func + + +def apply_rotary_emb_mageflow(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + """Apply complex rotary embeddings to `x` ([S, H, D]) using `freqs_cis` + (the MageFlowEmbedRope 2D multi-scale RoPE, adjacent-pair complex convention).""" + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(1) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(-2) + return x_out.type_as(x) + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +) -> torch.Tensor: + """Sinusoidal timestep embeddings (DDPM convention). + + NOTE: kept vendored (not diffusers') because the frequency table is + downcast to ``timesteps.dtype`` (bf16) here — the model was trained with + this exact bf16 rounding, so diffusers' fp32 variant produces a slightly + different embedding and degrades outputs. + """ + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timesteps.device + ) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent).to(timesteps.dtype) + emb = timesteps[:, None].float() * emb[None, :] + + # scale embeddings + emb = scale * emb + + # concat sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # flip sine and cosine embeddings + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # zero pad + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +class Timesteps(nn.Module): + def __init__( + self, + num_channels: int, + flip_sin_to_cos: bool, + downscale_freq_shift: float, + scale: int = 1, + ): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + return get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + ) + + +class MageFlowTimestepProjEmbeddings(nn.Module): + def __init__(self, embedding_dim): + super().__init__() + self.time_proj = Timesteps( + num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000 + ) + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim + ) + + def forward(self, timestep, hidden_states): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder( + timesteps_proj.to(dtype=hidden_states.dtype) + ) # (N, D) + return timesteps_emb + + +class MageFlowEmbedRope(nn.Module): + def __init__(self, theta: int, axes_dim: list, scale_rope=False): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + pos_index = torch.arange(4096) + neg_index = torch.arange(4096).flip(0) * -1 - 1 + self.pos_freqs = torch.cat( + [ + self.rope_params(pos_index, self.axes_dim[0], self.theta), + self.rope_params(pos_index, self.axes_dim[1], self.theta), + self.rope_params(pos_index, self.axes_dim[2], self.theta), + ], + dim=1, + ) + self.neg_freqs = torch.cat( + [ + self.rope_params(neg_index, self.axes_dim[0], self.theta), + self.rope_params(neg_index, self.axes_dim[1], self.theta), + self.rope_params(neg_index, self.axes_dim[2], self.theta), + ], + dim=1, + ) + + # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART + self.scale_rope = scale_rope + self.video_freq_cache = {} + + def rope_params(self, index, dim, theta=10000): + assert dim % 2 == 0 + freqs = torch.outer( + index, + 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)), + ) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + def forward( + self, video_fhw, device: torch.device, max_img_len: int = None + ) -> torch.Tensor: + """Compute the vision RoPE frequencies for the packed image tokens. + Text tokens are NOT rotated, so no text RoPE is computed. + + ``video_fhw`` is a list of (frame, height, width) tuples, one per image + segment in the packed sequence; segment index doubles as the "frame" + RoPE coordinate (edit reference images land on later frame indices). + """ + if self.pos_freqs.device != device: + self.pos_freqs = self.pos_freqs.to(device) + self.neg_freqs = self.neg_freqs.to(device) + + if isinstance(video_fhw, list): + video_fhw = video_fhw[0] + if not isinstance(video_fhw, list): + video_fhw = [video_fhw] + + vid_freqs = [] + for idx, fhw in enumerate(video_fhw): + frame, height, width = fhw + key = (frame, height, width, idx) + if key not in self.video_freq_cache: + self.video_freq_cache[key] = self._compute_video_freqs( + frame, height, width, idx + ) + vid_freqs.append(self.video_freq_cache[key].to(device)) + + vid_freqs = torch.cat(vid_freqs, dim=0) + + if max_img_len is not None and vid_freqs.shape[0] < max_img_len: + pad_len = max_img_len - vid_freqs.shape[0] + vid_freqs = torch.nn.functional.pad(vid_freqs, (0, 0, 0, pad_len)) + + return vid_freqs + + def _compute_video_freqs( + self, frame: int, height: int, width: int, idx: int = 0 + ) -> torch.Tensor: + seq_lens = frame * height * width + freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1) + freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1) + + freqs_frame = ( + freqs_pos[0][idx : idx + frame] + .view(frame, 1, 1, -1) + .expand(frame, height, width, -1) + ) + if self.scale_rope: + freqs_height = torch.cat( + [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], + dim=0, + ) + freqs_height = freqs_height.view(1, height, 1, -1).expand( + frame, height, width, -1 + ) + freqs_width = torch.cat( + [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], + dim=0, + ) + freqs_width = freqs_width.view(1, 1, width, -1).expand( + frame, height, width, -1 + ) + else: + freqs_height = ( + freqs_pos[1][:height] + .view(1, height, 1, -1) + .expand(frame, height, width, -1) + ) + freqs_width = ( + freqs_pos[2][:width] + .view(1, 1, width, -1) + .expand(frame, height, width, -1) + ) + + freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape( + seq_lens, -1 + ) + return freqs.clone().contiguous() + + +class Attention(nn.Module): + """Joint text+image attention (reference subset: the double-stream config + MageFlow actually instantiates — qk rms_norm, added kv projections).""" + + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + added_kv_proj_dim: int = None, + added_proj_bias: bool = True, + out_bias: bool = True, + eps: float = 1e-5, + processor=None, + out_dim: int = None, + out_context_dim: int = None, + elementwise_affine: bool = True, + ): + super().__init__() + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.inner_kv_dim = self.inner_dim + self.query_dim = query_dim + self.use_bias = bias + self.cross_attention_dim = query_dim + self.out_dim = out_dim if out_dim is not None else query_dim + self.out_context_dim = ( + out_context_dim if out_context_dim is not None else query_dim + ) + self.heads = out_dim // dim_head if out_dim is not None else heads + self.added_kv_proj_dim = added_kv_proj_dim + self.added_proj_bias = added_proj_bias + + # qk_norm is always "rms_norm" for MageFlow. + self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine) + + self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) + self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) + + self.add_k_proj = nn.Linear( + added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias + ) + self.add_v_proj = nn.Linear( + added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias + ) + self.add_q_proj = nn.Linear( + added_kv_proj_dim, self.inner_dim, bias=added_proj_bias + ) + self.norm_added_q = RMSNorm(dim_head, eps=eps) + self.norm_added_k = RMSNorm(dim_head, eps=eps) + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + self.to_out.append(nn.Dropout(dropout)) + + self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias) + + self.processor = processor + + def forward( + self, + hidden_states: torch.Tensor, + txt_cu_lens: torch.Tensor = None, + img_cu_lens: torch.Tensor = None, + image_rotary_emb: torch.Tensor = None, + **attention_kwargs, + ) -> torch.Tensor: + return self.processor( + self, + hidden_states, + txt_cu_lens=txt_cu_lens, + img_cu_lens=img_cu_lens, + image_rotary_emb=image_rotary_emb, + **attention_kwargs, + ) + + +class MageDoubleStreamAttnProcessor: + """Attention processor for the Mage double-stream architecture. Implements + joint attention where the packed text and image streams are concatenated + per sample (order: [text, image]) and run through one varlen kernel.""" + + def __call__( + self, + attn: Attention, + hidden_states: torch.FloatTensor, # Image stream [1, sum_img, D] + img_cu_lens: torch.LongTensor, + encoder_hidden_states: torch.FloatTensor = None, # Text stream [1, sum_txt, D] + txt_cu_lens: torch.LongTensor = None, + image_rotary_emb: torch.Tensor = None, + **kwargs, + ) -> torch.FloatTensor: + if encoder_hidden_states is None: + raise ValueError( + "MageDoubleStreamAttnProcessor requires encoder_hidden_states (text stream)" + ) + + # Compute QKV for image stream (sample projections) + img_query = attn.to_q(hidden_states) + img_key = attn.to_k(hidden_states) + img_value = attn.to_v(hidden_states) + + # Compute QKV for text stream (context projections) + txt_query = attn.add_q_proj(encoder_hidden_states) + txt_key = attn.add_k_proj(encoder_hidden_states) + txt_value = attn.add_v_proj(encoder_hidden_states) + + # Reshape for multi-head attention + img_query = img_query.unflatten(-1, (attn.heads, -1)) + img_key = img_key.unflatten(-1, (attn.heads, -1)) + img_value = img_value.unflatten(-1, (attn.heads, -1)) + + txt_query = txt_query.unflatten(-1, (attn.heads, -1)) + txt_key = txt_key.unflatten(-1, (attn.heads, -1)) + txt_value = txt_value.unflatten(-1, (attn.heads, -1)) + + if img_query.ndim == 4: + img_query = img_query.flatten(0, 1) + img_key = img_key.flatten(0, 1) + img_value = img_value.flatten(0, 1) + + if txt_query.ndim == 4: + txt_query = txt_query.flatten(0, 1) + txt_key = txt_key.flatten(0, 1) + txt_value = txt_value.flatten(0, 1) + + # Apply QK normalization + img_query = attn.norm_q(img_query) + img_key = attn.norm_k(img_key) + txt_query = attn.norm_added_q(txt_query) + txt_key = attn.norm_added_k(txt_key) + + # Apply 2D multi-scale RoPE (MageFlowEmbedRope) to image tokens + img_freqs = image_rotary_emb + img_query = apply_rotary_emb_mageflow(img_query, img_freqs) + img_key = apply_rotary_emb_mageflow(img_key, img_freqs) + + # Calculate lengths + img_lens = img_cu_lens[1:] - img_cu_lens[:-1] + txt_lens = txt_cu_lens[1:] - txt_cu_lens[:-1] + + # Calculate joint cu_seqlens + joint_lens = txt_lens + img_lens + joint_cu_lens = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=joint_lens.device), + torch.cumsum(joint_lens, dim=0, dtype=torch.int32), + ], + dim=0, + ) + + device = joint_lens.device + batch_size = len(txt_lens) + sample_indices = torch.arange(batch_size, device=device) + + txt_sample_ids = torch.repeat_interleave(sample_indices, txt_lens) + img_sample_ids = torch.repeat_interleave(sample_indices, img_lens) + + txt_intra_pos = ( + torch.arange(txt_query.shape[0], device=device) + - txt_cu_lens[txt_sample_ids] + ) + img_intra_pos = ( + torch.arange(img_query.shape[0], device=device) + - img_cu_lens[img_sample_ids] + ) + + txt_dest_indices = joint_cu_lens[txt_sample_ids] + txt_intra_pos + img_dest_indices = ( + joint_cu_lens[img_sample_ids] + txt_lens[img_sample_ids] + img_intra_pos + ) + + total_tokens = joint_cu_lens[-1] + joint_query = torch.empty( + (total_tokens, *txt_query.shape[1:]), dtype=txt_query.dtype, device=device + ) + joint_key = torch.empty( + (total_tokens, *txt_key.shape[1:]), dtype=txt_key.dtype, device=device + ) + joint_value = torch.empty( + (total_tokens, *txt_value.shape[1:]), dtype=txt_value.dtype, device=device + ) + + joint_query[txt_dest_indices] = txt_query + joint_query[img_dest_indices] = img_query + + joint_key[txt_dest_indices] = txt_key + joint_key[img_dest_indices] = img_key + + joint_value[txt_dest_indices] = txt_value + joint_value[img_dest_indices] = img_value + + max_seqlen = joint_lens.max().item() + joint_attn_output = flash_attn_varlen_func( + joint_query, + joint_key, + joint_value, + cu_seqlens_q=joint_cu_lens, + cu_seqlens_k=joint_cu_lens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + dropout_p=0.0, + softmax_scale=None, + causal=False, + ) + + txt_attn_output = joint_attn_output[txt_dest_indices] + img_attn_output = joint_attn_output[img_dest_indices] + + img_attn_output = img_attn_output.flatten(1, 2) # (N, H, D) -> (N, H*D) + img_attn_output = img_attn_output.to(joint_query.dtype) + + txt_attn_output = txt_attn_output.flatten(1, 2) # (N, H, D) -> (N, H*D) + txt_attn_output = txt_attn_output.to(joint_query.dtype) + + img_attn_output = attn.to_out[0](img_attn_output) + img_attn_output = attn.to_out[1](img_attn_output) # dropout + + txt_attn_output = attn.to_add_out(txt_attn_output) + txt_attn_output = txt_attn_output.view( + encoder_hidden_states.shape[0], + encoder_hidden_states.shape[1], + txt_attn_output.shape[-1], + ) + + return img_attn_output, txt_attn_output + + +class MageFlowTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + eps: float = 1e-6, + ): + super().__init__() + + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + + # Image processing modules + self.img_mod = nn.Sequential( + nn.SiLU(), + nn.Linear( + dim, 6 * dim, bias=True + ), # For scale, shift, gate for norm1 and norm2 + ) + self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.attn = Attention( + query_dim=dim, + added_kv_proj_dim=dim, # Enable added KV projections for text stream + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=dim, + bias=True, + processor=MageDoubleStreamAttnProcessor(), + eps=eps, + ) + self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_mlp = FeedForward( + dim=dim, dim_out=dim, activation_fn="gelu-approximate" + ) + + # Text processing modules + self.txt_mod = nn.Sequential( + nn.SiLU(), + nn.Linear( + dim, 6 * dim, bias=True + ), # For scale, shift, gate for norm1 and norm2 + ) + self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + # Text doesn't need separate attention - it's handled by img attn joint computation + self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_mlp = FeedForward( + dim=dim, dim_out=dim, activation_fn="gelu-approximate" + ) + + def _modulate(self, x, mod_params, cu_lens=None): + """Apply modulation to input tensor""" + shift, scale, gate = mod_params.chunk(3, dim=-1) + if cu_lens is not None: + assert x.shape[0] == 1, "x must be of shape (1, *) when cu_lens is not None" + x_flattened = x.view(-1, x.shape[-1]) + lengths = cu_lens[1:] - cu_lens[:-1] + shift_t = shift.repeat_interleave(lengths, dim=0) + scale_t = scale.repeat_interleave(lengths, dim=0) + gate_t = gate.repeat_interleave(lengths, dim=0) + + x_flattened = x_flattened * (1 + scale_t) + shift_t + x = x_flattened.view(x.shape) + return x, gate_t + else: + return x * (1 + scale) + shift, gate + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: torch.Tensor, + txt_cu_lens: torch.Tensor, + img_cu_lens: torch.Tensor, + joint_attention_kwargs: dict = None, + ): + # Get modulation parameters for both streams + img_mod_params = self.img_mod(temb) # [B, 6*dim] + txt_mod_params = self.txt_mod(temb) # [B, 6*dim] + + # Split modulation parameters for norm1 and norm2 + img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] + txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] + + # Process image stream - norm1 + modulation + img_normed = self.img_norm1(hidden_states) + img_modulated, img_gate1 = self._modulate( + img_normed, img_mod1, cu_lens=img_cu_lens + ) + + # Process text stream - norm1 + modulation + txt_normed = self.txt_norm1(encoder_hidden_states) + txt_modulated, txt_gate1 = self._modulate( + txt_normed, txt_mod1, cu_lens=txt_cu_lens + ) + + # Joint attention: computes QKV for both streams, applies QK norm and + # RoPE, concatenates per sample and splits the results back. + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states=img_modulated, + encoder_hidden_states=txt_modulated, + image_rotary_emb=image_rotary_emb, + txt_cu_lens=txt_cu_lens, + img_cu_lens=img_cu_lens, + **joint_attention_kwargs, + ) + + img_attn_output, txt_attn_output = attn_output + + # Apply attention gates and add residual + hidden_states = hidden_states + img_gate1 * img_attn_output + encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output + + # Process image stream - norm2 + MLP + img_normed2 = self.img_norm2(hidden_states) + img_modulated2, img_gate2 = self._modulate( + img_normed2, img_mod2, cu_lens=img_cu_lens + ) + img_mlp_output = self.img_mlp(img_modulated2) + hidden_states = hidden_states + img_gate2 * img_mlp_output + + # Process text stream - norm2 + MLP + txt_normed2 = self.txt_norm2(encoder_hidden_states) + txt_modulated2, txt_gate2 = self._modulate( + txt_normed2, txt_mod2, cu_lens=txt_cu_lens + ) + txt_mlp_output = self.txt_mlp(txt_modulated2) + encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output + + # Clip to prevent overflow for fp16 + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +class AdaLayerNormContinuous(nn.Module): + """Adaptive output norm; supports packed sequences via ``cu_seqlens``.""" + + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + elementwise_affine=True, + eps=1e-5, + bias=True, + norm_type="layer_norm", + ): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear( + conditioning_embedding_dim, embedding_dim * 2, bias=bias + ) + if norm_type == "layer_norm": + self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) + elif norm_type == "rms_norm": + self.norm = RMSNorm(embedding_dim, eps, elementwise_affine) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + def forward( + self, + x: torch.Tensor, + conditioning_embedding: torch.Tensor, + cu_seqlens: torch.Tensor = None, + ) -> torch.Tensor: + emb = self.linear(self.silu(conditioning_embedding).to(x.dtype)) + if cu_seqlens is None: + scale, shift = torch.chunk(emb, 2, dim=-1) + x = self.norm(x) * (1 + scale) + shift + else: + sample_lens = cu_seqlens[1:] - cu_seqlens[:-1] + flattened_x = x.view(-1, x.shape[-1]) + scale, shift = torch.chunk(emb, 2, dim=-1) + scale_t = torch.repeat_interleave(scale, sample_lens, dim=0) + shift_t = torch.repeat_interleave(shift, sample_lens, dim=0) + flattened_x = self.norm(flattened_x) * (1 + scale_t) + shift_t + x = flattened_x.view(x.shape) + return x + + +@dataclass +class MageFlowParams: + in_channels: int + out_channels: int + context_in_dim: int + hidden_size: int + num_heads: int + depth: int + axes_dim: list + checkpoint: bool + patch_size: int = 1 + + +class MageFlow(nn.Module): + def __init__(self, params: MageFlowParams): + super().__init__() + self.params = params + self.gradient_checkpointing = bool(params.checkpoint) + self.in_channels = params.in_channels + self.out_channels = params.out_channels + self.inner_dim = params.hidden_size # num_attention_heads * attention_head_dim + self.axes_dim = params.axes_dim + self.num_attention_heads = params.num_heads + self.attention_head_dim = self.inner_dim // self.num_attention_heads + self.patch_size = params.patch_size + assert sum(self.axes_dim) == self.attention_head_dim + + self.pos_embed = MageFlowEmbedRope( + theta=10000, axes_dim=self.axes_dim, scale_rope=True + ) + self.img_in = nn.Linear(self.in_channels, self.inner_dim) + self.txt_norm = RMSNorm(params.context_in_dim, eps=1e-6) + self.txt_in = nn.Linear(params.context_in_dim, self.inner_dim) + + self.time_text_embed = MageFlowTimestepProjEmbeddings( + embedding_dim=self.inner_dim + ) + + self.transformer_blocks = nn.ModuleList( + [ + MageFlowTransformerBlock( + dim=self.inner_dim, + num_attention_heads=self.num_attention_heads, + attention_head_dim=self.attention_head_dim, + ) + for _ in range(params.depth) + ] + ) + + self.norm_out = AdaLayerNormContinuous( + self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6 + ) + self.proj_out = nn.Linear( + self.inner_dim, + self.patch_size * self.patch_size * self.out_channels, + bias=True, + ) + + def enable_gradient_checkpointing(self, enable: bool = True): + self.gradient_checkpointing = enable + + def disable_gradient_checkpointing(self): + self.gradient_checkpointing = False + + def reset_rope(self): + """Rebuild the RoPE frequency tables. Needed after a meta-device init + (the tables are plain tensors, not buffers, so ``load_state_dict`` / + ``assign=True`` never materializes them).""" + self.pos_embed = MageFlowEmbedRope( + theta=10000, axes_dim=self.axes_dim, scale_rope=True + ) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def forward( + self, + img: Tensor, + txt: Tensor, + timesteps: Tensor, + img_shapes=None, + img_cu_seqlens: Tensor = None, + txt_cu_seqlens: Tensor = None, + attention_kwargs: dict = None, + ) -> Tensor: + if img.ndim != 3 or txt.ndim != 3: + raise ValueError("Input img and txt tensors must have 3 dimensions.") + + # Prepare vision RoPE (msrope); text tokens are not rotated. + ms_pe = self.pos_embed(img_shapes, device=img.device) + + img = self.img_in(img) + txt = self.txt_norm(txt) + + timesteps = timesteps.to(img.dtype) + temb = self.time_text_embed(timesteps, img) + + txt = self.txt_in(txt) + # The reference adds a zero "text vector" to temb (vec_type is null in + # the released config); kept for parity with the reference forward. + txt_vec = torch.zeros( + txt.shape[0], self.inner_dim, dtype=txt.dtype, device=txt.device + ) + temb = temb + txt_vec + + attention_kwargs = attention_kwargs or {} + + for block in self.transformer_blocks: + if self.gradient_checkpointing and torch.is_grad_enabled(): + txt, img = checkpoint( + block, + img, # hidden_states + txt, # encoder_hidden_states + temb, # temb + ms_pe, # image_rotary_emb + txt_cu_seqlens, # txt_cu_lens + img_cu_seqlens, # img_cu_lens + use_reentrant=False, + ) + else: + txt, img = block( + hidden_states=img, + encoder_hidden_states=txt, + txt_cu_lens=txt_cu_seqlens, + img_cu_lens=img_cu_seqlens, + temb=temb, + image_rotary_emb=ms_pe, + joint_attention_kwargs=attention_kwargs, + ) + + # Use only the image part (hidden_states) from the dual-stream blocks + img = self.norm_out(img, temb, cu_seqlens=img_cu_seqlens) + img = self.proj_out(img) + return img diff --git a/extensions_built_in/diffusion_models/mageflow/src/vae.py b/extensions_built_in/diffusion_models/mageflow/src/vae.py new file mode 100644 index 00000000..19953157 --- /dev/null +++ b/extensions_built_in/diffusion_models/mageflow/src/vae.py @@ -0,0 +1,703 @@ +""" +MageVAE: DConvEncoder + DConvDenoiser (with CoD Decoder) wrapper. + +Vendored from the reference implementation +(github.com/microsoft/Mage, mage_flow/models/modules/mage_vae.py); loguru +swapped for prints, otherwise unchanged. Loads the kl0.1 CoD ckpt layout used +by ``vae/diffusion_pytorch_model.safetensors`` in the Mage-Flow HF repos: + encoder weights: 'student.dconv_encoder.*' (packed mean+logvar, out_ch_mult=2) + decoder weights: 'pipeline.*' (denoiser + y_embedder.decoder) + +Latent shape: [B, 128, H/16, W/16] — no patch packing, no normalization. +Encode input / decode output are [-1, 1] pixel images with H, W multiples of 16. +""" + +import math +import os +from functools import lru_cache + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +# --------------------------------------------------------------------------- +# Primitive layers (vendored from GenCodec, inference subset) +# --------------------------------------------------------------------------- +def nonlinearity(x): + return x * torch.sigmoid(x) + + +def Normalize(in_channels): + return torch.nn.GroupNorm( + num_groups=32, num_channels=in_channels, eps=1e-6, affine=True + ) + + +def modulate(x, shift, scale): + if x.dim() == 4: + b, c = x.shape[:2] + return x * (1 + scale.view(b, c, 1, 1)) + shift.view(b, c, 1, 1) + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +class LayerNorm2d(nn.LayerNorm): + def __init__(self, num_channels, eps=1e-6, affine=True): + super().__init__(num_channels, eps=eps, elementwise_affine=affine) + + def forward(self, x): + # .contiguous() prevents a channels_last-strided NCHW view from + # propagating into downstream depthwise convs, which would otherwise + # hit a slow cuDNN path with a per-shape heuristic search. + x = x.permute(0, 2, 3, 1).contiguous() + x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + return x.permute(0, 3, 1, 2).contiguous() + + +class _EncoderLayerNorm2d(LayerNorm2d): + pass + + +class RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, x): + in_dtype = x.dtype + x = x.to(torch.float32) + var = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.variance_epsilon) + return self.weight * x.to(in_dtype) + + +class TimestepEmbedder(nn.Module): + """DConv-style timestep MLP (max_period=10000, freq_size=256, hidden=384).""" + + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half + ).to(t.device) + args = t[:, None].float() * freqs[None] + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1) + return emb + + def forward(self, t): + emb = self.timestep_embedding(t, self.frequency_embedding_size) + return self.mlp(emb.to(self.mlp[0].weight.dtype)) + + +class BottleneckPatchEmbed(nn.Module): + """Image patch embed concatenated with a per-patch conditioning vector.""" + + def __init__( + self, patch_size=16, in_chans=3, pca_dim=128, embed_dim=384, bias=True + ): + super().__init__() + self.proj1 = nn.Conv2d( + in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False + ) + self.proj2 = nn.Conv2d(pca_dim + embed_dim, embed_dim, kernel_size=1, bias=bias) + + def forward(self, x, cond): + return self.proj2(torch.cat([self.proj1(x), cond], dim=1)) + + +class DiCoBlock(nn.Module): + """DConv block with adaLN modulation.""" + + def __init__(self, hidden_size, mlp_ratio=4.0): + super().__init__() + self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True) + self.conv2 = nn.Conv2d( + hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True + ) + self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True) + + self.ca = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(hidden_size, hidden_size, 1, bias=True), + nn.Sigmoid(), + ) + + ffn = int(mlp_ratio * hidden_size) + self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True) + self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True) + + self.norm1 = LayerNorm2d(hidden_size, affine=False) + self.norm2 = LayerNorm2d(hidden_size, affine=False) + + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 6 * hidden_size, bias=True), + ) + + def forward(self, inp, c): + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.adaLN_modulation(c).chunk(6, dim=1) + ) + x = modulate(self.norm1(inp), shift_msa, scale_msa) + x = F.gelu(self.conv2(self.conv1(x))) + x = x * self.ca(x) + x = self.conv3(x) + x = inp + gate_msa[..., None, None] * x + x = x + gate_mlp[..., None, None] * self.conv5( + F.gelu(self.conv4(modulate(self.norm2(x), shift_mlp, scale_mlp))) + ) + return x + + +class _EncoderDiCoBlock(nn.Module): + """DiCoBlock without adaLN, for the encoder pathway.""" + + def __init__(self, hidden_size, mlp_ratio=4.0): + super().__init__() + self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True) + self.conv2 = nn.Conv2d( + hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True + ) + self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True) + self.ca = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(hidden_size, hidden_size, 1, bias=True), + nn.Sigmoid(), + ) + ffn = int(mlp_ratio * hidden_size) + self.conv4 = nn.Conv2d(hidden_size, ffn, 1, bias=True) + self.conv5 = nn.Conv2d(ffn, hidden_size, 1, bias=True) + self.norm1 = _EncoderLayerNorm2d(hidden_size) + self.norm2 = _EncoderLayerNorm2d(hidden_size) + + def forward(self, inp): + x = self.norm1(inp) + x = F.gelu(self.conv2(self.conv1(x))) + x = x * self.ca(x) + x = self.conv3(x) + x = inp + x + return x + self.conv5(F.gelu(self.conv4(self.norm2(x)))) + + +class NerfEmbedder(nn.Module): + """Patch-position embedder used by the DConv decoder x-pathway.""" + + def __init__(self, in_channels, hidden_size_input, max_freqs=8): + super().__init__() + self.max_freqs = max_freqs + self.embedder = nn.Sequential( + nn.Linear(in_channels + max_freqs**2, hidden_size_input, bias=True), + ) + + @lru_cache + def fetch_pos(self, patch_size, device, dtype): + pos = torch.linspace(0, 1, patch_size, device=device, dtype=dtype) + pos_y, pos_x = torch.meshgrid(pos, pos, indexing="ij") + pos_x = pos_x.reshape(-1, 1, 1) + pos_y = pos_y.reshape(-1, 1, 1) + freqs = torch.linspace( + 0, self.max_freqs, self.max_freqs, dtype=dtype, device=device + ) + fx = freqs[None, :, None] + fy = freqs[None, None, :] + coeffs = (1 + fx * fy) ** -1 + dct_x = torch.cos(pos_x * fx * torch.pi) + dct_y = torch.cos(pos_y * fy * torch.pi) + return (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs**2) + + def forward(self, x): + B, P2, _ = x.shape + ps = int(P2**0.5) + dct = self.fetch_pos(ps, x.device, x.dtype).expand(B, -1, -1) + return self.embedder(torch.cat([x, dct], dim=-1)) + + +class NerfFinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm = RMSNorm(hidden_size) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + + def forward(self, x): + return self.linear(self.norm(x)) + + +class SimpleMLPAdaLN(nn.Module): + """Final small MLP that maps NerfEmbedder features to per-patch RGB.""" + + def __init__( + self, + in_channels, + model_channels, + out_channels, + z_channels, + num_res_blocks, + patch_size, + ): + super().__init__() + self.in_channels = in_channels + self.model_channels = model_channels + self.out_channels = out_channels + self.num_res_blocks = num_res_blocks + self.patch_size = patch_size + + self.cond_embed = nn.Linear(z_channels, patch_size**2 * model_channels) + self.input_proj = nn.Linear(in_channels, model_channels) + + self.res_blocks = nn.ModuleList( + _MLPResBlock(model_channels) for _ in range(num_res_blocks) + ) + + def forward(self, x, c): + x = self.input_proj(x) + c = self.cond_embed(c).reshape(c.shape[0], self.patch_size**2, -1) + for block in self.res_blocks: + x = block(x, c) + return x + + +class _MLPResBlock(nn.Module): + def __init__(self, channels): + super().__init__() + self.in_ln = nn.LayerNorm(channels, eps=1e-6) + self.mlp = nn.Sequential( + nn.Linear(channels, channels, bias=True), + nn.SiLU(), + nn.Linear(channels, channels, bias=True), + ) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(channels, 3 * channels, bias=True), + ) + + def forward(self, x, y): + shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1) + h = self.in_ln(x) * (1 + scale) + shift + return x + gate * self.mlp(h) + + +class ResnetBlock(nn.Module): + """GroupNorm + Conv ResBlock used by the CoD Decoder.""" + + def __init__(self, *, in_channels, out_channels=None, dropout=0.0): + super().__init__() + out_channels = out_channels or in_channels + self.in_channels = in_channels + self.out_channels = out_channels + + self.norm1 = Normalize(in_channels) + self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1) + self.norm2 = Normalize(out_channels) + self.dropout = nn.Dropout(dropout) + self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1) + if in_channels != out_channels: + self.nin_shortcut = nn.Conv2d(in_channels, out_channels, 1) + + def forward(self, x): + h = self.conv1(nonlinearity(self.norm1(x))) + h = self.conv2(self.dropout(nonlinearity(self.norm2(h)))) + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class AttnBlock(nn.Module): + """Patched self-attention used at inference (eval mode of the original).""" + + def __init__(self, in_channels, patch_size=32): + super().__init__() + self.in_channels = in_channels + self.patch_size = patch_size + self.norm = Normalize(in_channels) + self.q = nn.Conv2d(in_channels, in_channels, 1) + self.k = nn.Conv2d(in_channels, in_channels, 1) + self.v = nn.Conv2d(in_channels, in_channels, 1) + self.proj_out = nn.Conv2d(in_channels, in_channels, 1) + + def forward(self, x): + h_ = self.norm(x) + Q = self.q(h_) + K = self.k(h_) + V = self.v(h_) + + d = self.patch_size + b, c, H, W = Q.shape + pad_h = (d - H % d) % d + pad_w = (d - W % d) % d + if pad_h or pad_w: + Q = F.pad(Q, (0, pad_w, 0, pad_h), mode="replicate") + K = F.pad(K, (0, pad_w, 0, pad_h), mode="replicate") + V = F.pad(V, (0, pad_w, 0, pad_h), mode="replicate") + _, _, H_pad, W_pad = Q.shape + nph, npw = H_pad // d, W_pad // d + np_ = nph * npw + + def to_patches(t): + return ( + t.reshape(b, c, nph, d, npw, d) + .permute(0, 2, 4, 1, 3, 5) + .reshape(b * np_, c, d * d) + ) + + Q = to_patches(Q) + K = to_patches(K) + V = to_patches(V) + + w_ = torch.bmm(Q.permute(0, 2, 1), K) * (c**-0.5) + w_ = F.softmax(w_, dim=2).permute(0, 2, 1) + h_ = ( + torch.bmm(V, w_) + .reshape(b, nph, npw, c, d, d) + .permute(0, 3, 1, 4, 2, 5) + .reshape(b, c, H_pad, W_pad) + ) + if pad_h or pad_w: + h_ = h_[:, :, :H, :W] + return x + self.proj_out(h_) + + +# --------------------------------------------------------------------------- +# adaLN constant-folding: at fixed t=0, adaLN_modulation(c) is constant. +# Replace the MLP with a buffer so DiCoBlock.forward stays unchanged. +# --------------------------------------------------------------------------- +class _ConstAdaLN(nn.Module): + def __init__(self, modulation: torch.Tensor): + super().__init__() + self.register_buffer("modulation", modulation.detach().clone()) + + def forward(self, c): + b = c.shape[0] + if self.modulation.shape[0] != b: + return self.modulation.expand(b, *self.modulation.shape[1:]) + return self.modulation + + +def _replace_adaln_with_const(module: nn.Module, c: torch.Tensor) -> int: + # Only DiCoBlock is targeted: its adaLN is conditioned solely on t. + # Other adaLN_modulation submodules (e.g. _MLPResBlock in the decoder MLP) + # take a per-position latent and must not be folded. + n = 0 + for child in module.modules(): + if not isinstance(child, DiCoBlock): + continue + adaln = child.adaLN_modulation + if isinstance(adaln, _ConstAdaLN): + continue + with torch.no_grad(): + mod = adaln(c) + child.adaLN_modulation = _ConstAdaLN(mod) + n += 1 + return n + + +# --------------------------------------------------------------------------- +# CoD Decoder: latent → conditioning features for the denoiser +# --------------------------------------------------------------------------- +class _Decoder(nn.Module): + """ds=16, up2x=True, light=True only.""" + + def __init__(self, out_ch=384, z_ch=128): + super().__init__() + self.conv_in = nn.Conv2d(z_ch, out_ch, kernel_size=3, stride=1, padding=1) + self.block = nn.Sequential( + ResnetBlock(in_channels=out_ch, out_channels=out_ch), + AttnBlock(out_ch, patch_size=32), + ResnetBlock(in_channels=out_ch, out_channels=out_ch), + AttnBlock(out_ch, patch_size=32), + ResnetBlock(in_channels=out_ch, out_channels=out_ch), + ) + self.norm_out = Normalize(out_ch) + self.conv_out = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1) + self.ada = nn.Identity() + + def forward(self, z): + h = self.block(self.conv_in(z)) + h = self.conv_out(nonlinearity(self.norm_out(h))) + return self.ada(h) + + +# --------------------------------------------------------------------------- +# DConvEncoder: image → packed (mean, logvar) latent +# --------------------------------------------------------------------------- +class _DConvEncoder(nn.Module): + def __init__( + self, + z_ch=128, + hidden_size=384, + num_blocks=21, + patch_size=16, + mlp_ratio=4.0, + head_size=768, + num_head_blocks=2, + out_ch_mult=2, + ): + super().__init__() + self.z_ch = z_ch + self.patch_size = patch_size + self.patch_cond_embed = nn.Conv2d( + 3, head_size, kernel_size=patch_size, stride=patch_size, bias=True + ) + self.head_blocks = nn.ModuleList( + [ + _EncoderDiCoBlock(head_size, mlp_ratio=mlp_ratio) + for _ in range(num_head_blocks) + ] + ) + self.proj_down = nn.Conv2d(head_size, hidden_size, kernel_size=1, bias=True) + self.z_proj = nn.Conv2d(z_ch, hidden_size, kernel_size=1, bias=True) + self.fuse_proj = nn.Conv2d( + hidden_size * 2, hidden_size, kernel_size=1, bias=True + ) + self.t_embedder = TimestepEmbedder(hidden_size) + self.blocks = nn.ModuleList( + [DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks)] + ) + self.norm_out = LayerNorm2d(hidden_size) + self.proj_out = nn.Conv2d( + hidden_size, z_ch * out_ch_mult, kernel_size=1, bias=True + ) + + def forward_pred(self, z_t, t, y): + cond = self.patch_cond_embed(y) + for block in self.head_blocks: + cond = block(cond) + cond = self.proj_down(cond) + + s = self.fuse_proj(torch.cat([cond, self.z_proj(z_t)], dim=1)) + c = self.t_embedder(t.view(-1)) + for block in self.blocks: + s = block(s, c) + return self.proj_out(self.norm_out(s)) + + +# --------------------------------------------------------------------------- +# DConv denoiser: latent (via cond) + zero noise → reconstructed image +# --------------------------------------------------------------------------- +class _YEmbedder(nn.Module): + """Holds only the CoD decoder; the original Flux2 VAE encoder side is omitted.""" + + def __init__(self, ch=384, z_ch=128): + super().__init__() + self.decoder = _Decoder(out_ch=ch, z_ch=z_ch) + + +class _DConvDenoiser(nn.Module): + def __init__( + self, + patch_size=16, + in_channels=3, + hidden_size=384, + hidden_size_x=32, + mlp_ratio=4.0, + num_blocks=24, + num_cond_blocks=21, + bottleneck_dim=128, + ): + super().__init__() + self.in_channels = in_channels + self.patch_size = patch_size + self.hidden_size = hidden_size + self.num_cond_blocks = num_cond_blocks + + self.t_embedder = TimestepEmbedder(hidden_size) + self.y_embedder_x = nn.Conv2d( + hidden_size, hidden_size_x * patch_size**2, 1, 1, 0 + ) + self.x_embedder = NerfEmbedder( + in_channels + hidden_size_x, hidden_size_x, max_freqs=8 + ) + self.s_embedder = BottleneckPatchEmbed( + patch_size, in_channels, bottleneck_dim, hidden_size, bias=True + ) + self.blocks = nn.ModuleList( + [ + DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) + for _ in range(num_cond_blocks) + ] + ) + self.dec_net = SimpleMLPAdaLN( + in_channels=hidden_size_x, + model_channels=hidden_size_x, + out_channels=in_channels, + z_channels=hidden_size, + num_res_blocks=num_blocks - num_cond_blocks, + patch_size=patch_size, + ) + self.final_layer = NerfFinalLayer(hidden_size_x, in_channels) + self.y_embedder = _YEmbedder(ch=hidden_size, z_ch=bottleneck_dim) + + def forward(self, x, t, cond): + b, _, h, w = x.shape + c = self.t_embedder(t.view(-1)) + + s = self.s_embedder(x, cond) + for block in self.blocks: + s = block(s, c) + + length = s.shape[-2] * s.shape[-1] + s = s.permute(0, 2, 3, 1).reshape(-1, self.hidden_size) + + x = torch.nn.functional.unfold( + x, kernel_size=self.patch_size, stride=self.patch_size + ) + x = torch.cat([x, self.y_embedder_x(cond).flatten(2)], dim=1) + x = ( + x.reshape(b, -1, self.patch_size**2, length) + .permute(0, 3, 2, 1) + .flatten(0, 1) + ) + x = self.x_embedder(x) + + x = self.dec_net(x, s) + x = self.final_layer(x) + x = x.transpose(1, 2).reshape(b, length, -1) + return torch.nn.functional.fold( + x.transpose(1, 2).contiguous(), + (h, w), + kernel_size=self.patch_size, + stride=self.patch_size, + ) + + +# --------------------------------------------------------------------------- +# Wrapper +# --------------------------------------------------------------------------- +def _load_state_dict(ckpt_path: str): + if ckpt_path.endswith(".safetensors"): + from safetensors.torch import load_file + + return load_file(ckpt_path, device="cpu") + if os.path.exists(os.path.join(ckpt_path, "checkpoint-state_dict.pt")): + ckpt_path = os.path.join(ckpt_path, "checkpoint-state_dict.pt") + elif os.path.isdir(ckpt_path): + ckpt_path = os.path.join(ckpt_path, "checkpoint", "mp_rank_00_model_states.pt") + state = torch.load(ckpt_path, map_location="cpu") + if "module" in state: + return state["module"] + if "state_dict" in state: + return state["state_dict"] + return state + + +class MageVAE(nn.Module): + """ + Encode: DConvEncoder (one-step diffusion) → latent [B, 128, H/16, W/16] + Decode: DConvDenoiser + CoD Decoder → image [B, 3, H, W] in [-1, 1] + """ + + latent_channels = 128 + downsample_factor = 16 + + def __init__(self, ckpt_path: str, sample_posterior: bool = True): + super().__init__() + self.sample_posterior = sample_posterior + + self.dconv_encoder = _DConvEncoder() + self.decoder_model = _DConvDenoiser() + + sd = _load_state_dict(ckpt_path) + self._load_encoder(sd, ckpt_path) + self._load_decoder(sd, ckpt_path) + + # adaLN modulation depends only on t, and we always run at t=0. + # Precompute and drop the MLPs once at construction (~37M params saved). + self._freeze_adaln_cache() + + def _load_encoder(self, sd, ckpt_path): + prefix = "student.dconv_encoder." + enc_sd = {k[len(prefix) :]: v for k, v in sd.items() if k.startswith(prefix)} + if not enc_sd: + raise RuntimeError(f"MageVAE encoder: no '{prefix}*' keys in {ckpt_path}") + proj = enc_sd.get("proj_out.weight") + if proj is None or proj.shape[0] != 2 * self.latent_channels: + raise RuntimeError( + f"MageVAE encoder: expected packed mean+logvar (proj_out out_channels=" + f"{2 * self.latent_channels}), got {None if proj is None else tuple(proj.shape)}" + ) + missing, unexpected = self.dconv_encoder.load_state_dict(enc_sd, strict=False) + if missing: + print(f"MageVAE encoder missing keys: {missing[:10]}") + + def _load_decoder(self, sd, ckpt_path): + prefix = "pipeline." + if not any(k.startswith(prefix) for k in sd): + raise RuntimeError(f"MageVAE decoder: no '{prefix}*' keys in {ckpt_path}") + model_dict = self.decoder_model.state_dict() + matched = {} + for k, v in sd.items(): + if not k.startswith(prefix): + continue + new_k = k[len(prefix) :] + if new_k.startswith("y_embedder.encoder.") or new_k.startswith( + "y_embedder.bottleneck." + ): + continue + if new_k in model_dict and model_dict[new_k].shape == v.shape: + matched[new_k] = v + self.decoder_model.load_state_dict(matched, strict=False) + if not matched: + raise RuntimeError(f"MageVAE decoder: 0 params matched from {ckpt_path}") + + @torch.no_grad() + def _moments(self, x: torch.Tensor): + B, _, H, W = x.shape + ps = self.dconv_encoder.patch_size + z_t = torch.zeros( + B, self.dconv_encoder.z_ch, H // ps, W // ps, device=x.device, dtype=x.dtype + ) + t = torch.zeros(B, device=x.device, dtype=x.dtype) + out = self.dconv_encoder.forward_pred(z_t, t, x) + mean = out[:, : self.latent_channels] + logvar = out[:, self.latent_channels :].clamp(min=-20.0, max=10.0) + return mean, logvar + + @torch.no_grad() + def encode(self, x: torch.Tensor) -> torch.Tensor: + ps = self.dconv_encoder.patch_size + H, W = x.shape[-2], x.shape[-1] + if H % ps or W % ps: + raise ValueError(f"H, W must be multiples of {ps}, got ({H}, {W})") + mean, logvar = self._moments(x) + if self.sample_posterior: + return mean + torch.exp(0.5 * logvar) * torch.randn_like(mean) + return mean + + @torch.no_grad() + def decode(self, z: torch.Tensor) -> torch.Tensor: + cond = self.decoder_model.y_embedder.decoder(z) + B = z.shape[0] + H = z.shape[2] * self.downsample_factor + W = z.shape[3] * self.downsample_factor + noise = torch.zeros(B, 3, H, W, device=z.device, dtype=z.dtype) + t = torch.zeros(B, device=z.device, dtype=z.dtype) + return self.decoder_model.forward(noise, t, cond) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def _freeze_adaln_cache(self): + """Constant-fold adaLN_modulation MLPs at t=0 (encoder + decoder).""" + device = next(self.parameters()).device + dtype = next(self.parameters()).dtype + t = torch.zeros(1, device=device, dtype=dtype) + c_enc = self.dconv_encoder.t_embedder(t) + _replace_adaln_with_const(self.dconv_encoder, c_enc) + c_dec = self.decoder_model.t_embedder(t) + _replace_adaln_with_const(self.decoder_model, c_dec) diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts index 51dfb349..fd7ad830 100644 --- a/ui/src/app/jobs/new/options.ts +++ b/ui/src/app/jobs/new/options.ts @@ -1196,6 +1196,55 @@ export const modelArchs: ModelArch[] = [ 'model.model_kwargs.kv_cache', ], }, + { + name: 'mageflow', + label: 'Mage-Flow', + group: 'image', + defaults: { + 'config.process[0].model.name_or_path': ['microsoft/Mage-Flow-Base', defaultNameOrPath], + 'config.process[0].model.quantize': [true, false], + 'config.process[0].model.quantize_te': [true, false], + 'config.process[0].train.timestep_type': ['linear', 'sigmoid'], + 'config.process[0].network.conv': [undefined, 16], + 'config.process[0].network.conv_alpha': [undefined, 16], + 'config.process[0].model.low_vram': [true, false], + 'config.process[0].sample.guidance_scale': [4, 4], + 'config.process[0].sample.sample_steps': [25, 25], + }, + disableSections: [ + 'network.conv', + ], + additionalSections: [ + 'model.low_vram', + 'model.layer_offloading', + ], + }, + { + name: 'mageflow_edit', + label: 'Mage-Flow Edit', + group: 'instruction', + defaults: { + 'config.process[0].model.name_or_path': ['microsoft/Mage-Flow-Edit-Base', defaultNameOrPath], + 'config.process[0].model.quantize': [true, false], + 'config.process[0].model.quantize_te': [true, false], + 'config.process[0].train.timestep_type': ['linear', 'sigmoid'], + 'config.process[0].network.conv': [undefined, 16], + 'config.process[0].network.conv_alpha': [undefined, 16], + 'config.process[0].model.low_vram': [true, false], + 'config.process[0].sample.guidance_scale': [4, 4], + 'config.process[0].sample.sample_steps': [25, 25], + 'config.process[0].train.unload_text_encoder': [false, false], + }, + disableSections: [ + 'network.conv', 'train.unload_text_encoder', + ], + additionalSections: [ + 'datasets.multi_control_paths', + 'sample.multi_ctrl_imgs', + 'model.low_vram', + 'model.layer_offloading', + ], + }, { name: 'boogu_image', label: 'Boogu Image',