Merge pull request #1025 from ostris/models/v2

Reworked the entire model loading to make everything testable, hot swappable, and uniform.

.gitignore

@@ -191,4 +191,5 @@ aitk_db.db-shm
 /data
 .claude
 original_repo
-.next
+.next
+testing/.model_test_outputs

extensions_built_in/audio_models/ace_step/ace_step_15_model.py

@@ -11,9 +11,7 @@ from toolkit.prompt_utils import PromptEmbeds, concat_prompt_embeds
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
-from toolkit.util.quantize import get_qtype, quantize, quantize_model
 
-from optimum.quanto import freeze
 from .src.model import (
     AceStep15,
     OobleckVAE,
@@ -116,32 +114,22 @@ class AceStep15Model(BaseAudioModel):
         models = load_models(model_path, device=load_device, dtype=dtype)
 
         self.model = models["model"]
-        
-        if self.model_config.quantize:
-            self.print_and_status_update("Quantizing Transformer")
-            # quantize_model(self, self.model.decoder)
-            quantize(self.model, weights=get_qtype(self.model_config.qtype))
-            freeze(self.model)
-            flush()
-        
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            self.model.to("cpu")
-            
-        
+
         if (
             self.model_config.layer_offloading
             and self.model_config.layer_offloading_transformer_percent > 0
         ):
             raise NotImplementedError("Layer offloading not yet implemented for AceStep15Model")
-        
+
+        # quantize + offload + placement, all driven by model_config
+        self.model.aitk_post_load(**self.component_load_kwargs("transformer"))
+        flush()
+
         self.text_encoder = models["text_encoder"]
-        
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            quantize(self.text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(self.text_encoder)
-            flush()
+
+        # quantize + offload + placement, all driven by model_config
+        self.text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
+        flush()
         
         self.vae = models["vae"]
         

extensions_built_in/audio_models/ace_step/src/model.py

@@ -19,6 +19,8 @@ import os
 import time
 
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 import torch.nn.functional as F
 import torchaudio
 from safetensors.torch import load_file
@@ -857,7 +859,35 @@ class DiTModel(nn.Module):
 # ── Top-level model ──
 
 
-class AceStep15(nn.Module):
+class AceStep15(nn.Module, OstrisModelMixin):
+    @classmethod
+    def load_from_state_dict(cls, state_dict, dtype=None, **kwargs):
+        cfg = infer_dit_config(state_dict)
+        model = cls(
+            hidden=cfg["hidden"],
+            inter=cfg["inter"],
+            heads=cfg["heads"],
+            kv=cfg["kv"],
+            head_dim=cfg["head_dim"],
+            n_dit=cfg["n_dit"],
+            n_lyric=cfg["n_lyric"],
+            n_timbre=cfg["n_timbre"],
+            enc_hidden=cfg["enc_hidden"],
+            enc_heads=cfg["enc_heads"],
+            enc_kv=cfg["enc_kv"],
+            enc_inter=cfg["enc_inter"],
+        )
+        missing, unexpected = model.load_state_dict(state_dict, strict=False)
+        # tokenizer/detokenizer keys are expected to be unused (cover mode only)
+        unexpected = [
+            k for k in unexpected if not k.startswith(("tokenizer.", "detokenizer."))
+        ]
+        if missing:
+            print(f"    DiT missing: {len(missing)} (first 3: {missing[:3]})")
+        if unexpected:
+            print(f"    DiT unexpected: {len(unexpected)} (first 3: {unexpected[:3]})")
+        return model
+
     def __init__(
         self,
         hidden=2048,
@@ -1051,7 +1081,17 @@ class _SeqWrap(nn.Module):
         return self.layers(x)
 
 
-class OobleckVAE(nn.Module):
+class OobleckVAE(nn.Module, OstrisModelMixin):
+    @classmethod
+    def load_from_state_dict(cls, state_dict, dtype=None, **kwargs):
+        vae = cls()
+        m, u = vae.load_state_dict(state_dict, strict=False)
+        if m:
+            print(f"    VAE missing: {len(m)} (first 3: {m[:3]})")
+        if u:
+            print(f"    VAE unexpected: {len(u)}")
+        return vae
+
     def __init__(
         self,
         in_ch=2,
@@ -1152,7 +1192,28 @@ class OobleckVAE(nn.Module):
 # ═══════════════════════════════════════════════════════════════════════════════
 
 
-class TextEncoder(nn.Module):
+class TextEncoder(nn.Module, OstrisModelMixin):
+    @classmethod
+    def load_from_state_dict(cls, state_dict, dtype=None, **kwargs):
+        from transformers import Qwen3Config, Qwen3Model
+
+        qwen_cfg = Qwen3Config(
+            vocab_size=151669,
+            hidden_size=1024,
+            intermediate_size=3072,
+            num_hidden_layers=28,
+            num_attention_heads=16,
+            num_key_value_heads=8,
+            head_dim=128,
+            max_position_embeddings=32768,
+            rms_norm_eps=1e-6,
+        )
+        qwen = Qwen3Model(qwen_cfg)
+        m2, u2 = qwen.load_state_dict(state_dict, strict=False)
+        if m2:
+            print(f"    TE missing: {len(m2)} (first 3: {m2[:3]})")
+        return cls(qwen)
+
     """Wraps Qwen3 weights loaded from AIO. Forward returns last_hidden_state."""
 
     def __init__(self, qwen_model):
@@ -1261,41 +1322,13 @@ def load_models(checkpoint_path, device="cuda", dtype=torch.bfloat16):
         for k, v in sd.items()
         if k.startswith("model.diffusion_model.")
     }
-    cfg = infer_dit_config(dit_sd)
-    model = AceStep15(
-        hidden=cfg["hidden"],
-        inter=cfg["inter"],
-        heads=cfg["heads"],
-        kv=cfg["kv"],
-        head_dim=cfg["head_dim"],
-        n_dit=cfg["n_dit"],
-        n_lyric=cfg["n_lyric"],
-        n_timbre=cfg["n_timbre"],
-        enc_hidden=cfg["enc_hidden"],
-        enc_heads=cfg["enc_heads"],
-        enc_kv=cfg["enc_kv"],
-        enc_inter=cfg["enc_inter"],
-    )
-    missing, unexpected = model.load_state_dict(dit_sd, strict=False)
-    # tokenizer/detokenizer keys are expected to be unused (cover mode only)
-    unexpected = [
-        k for k in unexpected if not k.startswith(("tokenizer.", "detokenizer."))
-    ]
-    if missing:
-        print(f"    DiT missing: {len(missing)} (first 3: {missing[:3]})")
-    if unexpected:
-        print(f"    DiT unexpected: {len(unexpected)} (first 3: {unexpected[:3]})")
+    model = AceStep15.load_from_state_dict(dit_sd, dtype)
     model = model.to(device).to(dtype).eval()
 
     # --- VAE ---
     print("  Loading VAE...")
     vae_sd = {k.removeprefix("vae."): v for k, v in sd.items() if k.startswith("vae.")}
-    vae = OobleckVAE()
-    m, u = vae.load_state_dict(vae_sd, strict=False)
-    if m:
-        print(f"    VAE missing: {len(m)} (first 3: {m[:3]})")
-    if u:
-        print(f"    VAE unexpected: {len(u)}")
+    vae = OobleckVAE.load_from_state_dict(vae_sd, dtype)
     vae = vae.to(device).to(dtype).eval()
 
     # --- Text encoder (Qwen3-Embedding from AIO) ---
@@ -1305,25 +1338,7 @@ def load_models(checkpoint_path, device="cuda", dtype=torch.bfloat16):
         for k, v in sd.items()
         if k.startswith("text_encoders.qwen3_06b.transformer.model.")
     }
-    # Load Qwen3 model structure from transformers, then override weights
-    from transformers import Qwen3Model, Qwen3Config
-
-    qwen_cfg = Qwen3Config(
-        vocab_size=151669,
-        hidden_size=1024,
-        intermediate_size=3072,
-        num_hidden_layers=28,
-        num_attention_heads=16,
-        num_key_value_heads=8,
-        head_dim=128,
-        max_position_embeddings=32768,
-        rms_norm_eps=1e-6,
-    )
-    qwen = Qwen3Model(qwen_cfg)
-    m2, u2 = qwen.load_state_dict(te_sd, strict=False)
-    if m2:
-        print(f"    TE missing: {len(m2)} (first 3: {m2[:3]})")
-    te = TextEncoder(qwen).to(device).to(dtype).eval()
+    te = TextEncoder.load_from_state_dict(te_sd, dtype).to(device).to(dtype).eval()
 
     # Tokenizer — download from HF
     print("  Loading tokenizer...")

extensions_built_in/audio_models/base_audio_model.py

@@ -69,21 +69,7 @@ class BaseAudioModel(BaseModel):
             "save_model is not implemented for this model. Use the pipeline directly instead."
         )
 
-    def convert_lora_weights_before_save(self, state_dict):
-        # currently starte with transformer. but needs to start with diffusion_model. for comfyui
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
-
-    def convert_lora_weights_before_load(self, state_dict):
-        # saved as diffusion_model. but needs to be transformer. for ai-toolkit
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
     def encode_images(self, image_list: torch.Tensor, device=None, dtype=None):
         # make it more obvious for audio models

extensions_built_in/captioner/Qwen3VLCaptioner.py

@@ -12,6 +12,8 @@ from optimum.quanto import freeze
 from toolkit.basic import flush
 from toolkit.util.quantize import quantize, get_qtype
 
+from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed
+
 from .BaseCaptioner import BaseCaptioner
 import transformers
 import logging
@@ -19,29 +21,6 @@ import traceback
 import warnings
 
 
-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."""
-    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
-
-
 # transformers.logging.set_verbosity_error()
 warnings.filterwarnings("ignore")
 logging.disable(logging.WARNING)

extensions_built_in/captioner/prompts/ideogram4_prompt.py

@@ -1,4 +1,4 @@
-ideogram4_prompt = """
+ideogram4_prompt = r"""
 [META]
 frozen: false
 description: Slim single-shot magic prompt — splatter planning + v15 output discipline, deduped for faster inference. Thinking off.
@@ -15,7 +15,7 @@ You convert a natural-language user idea into a structured JSON caption an image
 
 - Emit a SINGLE-LINE MINIFIED JSON object — no markdown fences, no commentary, no other top-level keys.
 - Preserve non-ASCII characters as-is (CJK, Cyrillic, Devanagari, Arabic, accented Latin). Never escape with `\uNNNN`, transliterate, or replace `café` with `cafe`.
-- Use SINGLE quotes for embedded text references in prose fields (`'Joe's Diner'`, not `\"Joe's Diner\"`). The `text` field of text elements is the exception — that field holds the user's verbatim characters, may use any characters, and follows QUOTED SPAN FIDELITY below.
+- Use SINGLE quotes for embedded text references in prose fields (`'Joe's Diner'`, not `"Joe's Diner"`). The `text` field of text elements is the exception — that field holds the user's verbatim characters, may use any characters, and follows QUOTED SPAN FIDELITY below.
 
 ### Target aspect ratio (input only — never emit it)
 

extensions_built_in/diffusion_models/anima/anima.py

@@ -3,21 +3,21 @@ from typing import List, Optional
 
 import torch
 import yaml
-from optimum.quanto import freeze
 from safetensors.torch import load_file, save_file
 
 from toolkit.accelerator import unwrap_model
 from toolkit.basic import flush
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
-from toolkit.memory_management import MemoryManager
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.diffusion_models.cosmos import CosmosTransformer3DModel
+from toolkit.models.v2.text_encoders.anima import AnimaTextConditioner
+from toolkit.models.v2.text_encoders.qwen3 import Qwen3ModelEncoder
+from toolkit.models.v2.vae.qwen_image import QwenImageVAE
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
-from toolkit.util.quantize import get_qtype, quantize, quantize_model
 
 try:
-    from diffusers import AnimaAutoBlocks, AnimaModularPipeline, AnimaTextConditioner
-    from diffusers.models import CosmosTransformer3DModel
+    from diffusers import AnimaAutoBlocks, AnimaModularPipeline
     from diffusers.modular_pipelines import SequentialPipelineBlocks
     from diffusers.modular_pipelines.anima.modular_blocks_anima import AnimaCoreDenoiseStep, AnimaDecodeStep
 except ImportError as e:
@@ -255,72 +255,52 @@ class AnimaModel(BaseModel):
         self.print_and_status_update("Loading Anima model")
 
         pipe: AnimaModularPipeline = AnimaAutoBlocks().init_pipeline(self.model_config.name_or_path)
-        load_kwargs = {"torch_dtype": dtype}
-        model_path = os.path.abspath(os.path.expanduser(str(self.model_config.name_or_path)))
-        if os.path.isdir(model_path):
-            load_kwargs["pretrained_model_name_or_path"] = model_path
-        pipe.load_components(**load_kwargs)
-        pipe.update_components(scheduler=self.get_train_scheduler())
+        name = self.model_config.name_or_path
+        local_path = os.path.abspath(os.path.expanduser(str(name)))
+        if os.path.isdir(local_path):
+            name = local_path
+
+        # components load individually through the v2 module classes and are
+        # handed to the modular pipeline
+        from transformers import AutoTokenizer
+
+        self.print_and_status_update("Loading components")
+        transformer = CosmosTransformer3DModel.load_model(name, dtype=dtype)
+        vae = QwenImageVAE.load_model(name, dtype=dtype)
+        text_encoder = Qwen3ModelEncoder.load_model(name, dtype=dtype)
+        text_conditioner = AnimaTextConditioner.load_model(name, dtype=dtype)
+        tokenizer = AutoTokenizer.from_pretrained(name, subfolder="tokenizer")
+        t5_tokenizer = AutoTokenizer.from_pretrained(name, subfolder="t5_tokenizer")
+        pipe.update_components(
+            transformer=transformer,
+            vae=vae,
+            text_encoder=text_encoder,
+            text_conditioner=text_conditioner,
+            tokenizer=tokenizer,
+            t5_tokenizer=t5_tokenizer,
+            scheduler=self.get_train_scheduler(),
+        )
 
         transformer = pipe.transformer
         text_conditioner = pipe.text_conditioner
 
-        if self.model_config.quantize:
-            self.print_and_status_update("Quantizing Transformer")
-            quantize_model(self, transformer)
-            flush()
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
 
-            self.print_and_status_update("Quantizing Text Conditioner")
-            quantize(text_conditioner, weights=get_qtype(self.model_config.qtype))
-            freeze(text_conditioner)
-            flush()
+        # the text conditioner rides the transformer quantize flag (at qtype_te)
+        # but takes the text-encoder offload/placement policy
+        tc_kwargs = self.component_load_kwargs("te")
+        tc_kwargs["qtype"] = (
+            self.model_config.qtype_te if self.model_config.quantize else None
+        )
+        text_conditioner.aitk_post_load(**tc_kwargs)
+        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.layer_offloading
-            and self.model_config.layer_offloading_text_encoder_percent > 0
-        ):
-            MemoryManager.attach(
-                pipe.text_encoder,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_text_encoder_percent,
-            )
-            MemoryManager.attach(
-                text_conditioner,
-                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 transformer to CPU")
-            transformer.to("cpu")
-            text_conditioner.to("cpu")
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
-            text_conditioner.to(self.device_torch, dtype=dtype)
-
-        pipe.text_encoder.to(self.device_torch, dtype=dtype)
+        # quantize + offload + placement, all driven by model_config
+        pipe.text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
         pipe.text_encoder.requires_grad_(False)
         pipe.text_encoder.eval()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            quantize(pipe.text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(pipe.text_encoder)
-            flush()
-
-        if self.model_config.low_vram:
-            pipe.text_encoder.to("cpu")
-            flush()
+        flush()
 
         self.noise_scheduler = pipe.scheduler
         self.vae = pipe.vae

extensions_built_in/diffusion_models/boogu_image/boogu_image.py

@@ -30,12 +30,15 @@ from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
 from toolkit.basic import flush
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLModelEncoder
+from toolkit.models.v2.vae.autoencoder_kl import KLVAE
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
 
-from optimum.quanto import freeze, QTensor
+
+from optimum.quanto import QTensor
 from diffusers import AutoencoderKL
 
 from .src.transformer import BooguImageTransformer2DModel
@@ -71,33 +74,6 @@ SYSTEM_PROMPT_T2I = (
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 
 
-def patch_qwen_vl_patch_embed(model) -> int:
-    """Swap Qwen-VL's vision ``patch_embed`` Conv3d for the equivalent ``F.linear``.
-
-    Qwen-VL's patch_embed is a Conv3d whose kernel == stride, i.e. just a linear
-    projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and
-    falls back to a slow path that effectively locks up image caching for the edit
-    (TI2I) model. The weight is read lazily so this survives later ``.to()`` moves.
-    Returns the number of patch_embed modules patched. (Vendored from
-    extensions_built_in/captioner/Qwen3VLCaptioner.py.)
-    """
-    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
-
-
 class BooguImageModel(BaseModel):
     arch = "boogu_image"
     # Default HF repo when model.name_or_path is unset (overridden by the edit model).
@@ -167,8 +143,8 @@ class BooguImageModel(BaseModel):
         # deserialize -- use the bf16 repo and let ai-toolkit quantize if wanted.
         self.print_and_status_update("Loading transformer")
         try:
-            transformer = BooguImageTransformer2DModel.from_pretrained(
-                base, subfolder="transformer", torch_dtype=dtype, token=HF_TOKEN
+            transformer = BooguImageTransformer2DModel.load_model(
+                base, dtype=dtype, token=HF_TOKEN
             )
         except OSError as e:
             raise OSError(
@@ -188,16 +164,8 @@ class BooguImageModel(BaseModel):
         if attention_backend != "native":
             transformer.set_attention_backend(attention_backend)
 
-        if self.model_config.quantize:
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            flush()
-
-        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)
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
 
         # --- instruction encoder (Qwen3-VL) + processor ---
@@ -212,8 +180,8 @@ class BooguImageModel(BaseModel):
         # AutoModel yields the inner Qwen3VLModel (the ``.model`` of the
         # *ForConditionalGeneration), whose last_hidden_state is exactly the
         # instruction feature the Boogu pipeline consumes.
-        text_encoder = AutoModel.from_pretrained(
-            te_path, subfolder=te_subfolder, torch_dtype=dtype, token=HF_TOKEN
+        text_encoder = Qwen3VLModelEncoder.load_model(
+            te_path, dtype=dtype, subfolder=te_subfolder, token=HF_TOKEN
         )
         text_encoder.eval()
         text_encoder.requires_grad_(False)
@@ -227,25 +195,13 @@ class BooguImageModel(BaseModel):
             )
         flush()
 
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing instruction 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.low_vram:
-            self.print_and_status_update("Moving instruction encoder to CPU")
-            text_encoder.to("cpu")
-        else:
-            text_encoder.to(self.device_torch)
+        # quantize + offload + placement, all driven by model_config
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
         flush()
 
         # --- VAE (FLUX AutoencoderKL) ---
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKL.from_pretrained(
-            base, subfolder="vae", torch_dtype=self.vae_torch_dtype, token=HF_TOKEN
-        )
+        vae = KLVAE.load_model(base, dtype=self.vae_torch_dtype, token=HF_TOKEN)
         vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype)
         vae.eval()
         vae.requires_grad_(False)
@@ -446,14 +402,5 @@ class BooguImageModel(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["double_stream_layers", "single_stream_layers"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        return {
-            k.replace("transformer.", "diffusion_model."): v
-            for k, v in state_dict.items()
-        }
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        return {
-            k.replace("diffusion_model.", "transformer."): v
-            for k, v in state_dict.items()
-        }

extensions_built_in/diffusion_models/boogu_image/src/transformer.py

@@ -21,6 +21,8 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 import torch.nn as nn
 from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 from diffusers.loaders import PeftAdapterMixin
 from diffusers.loaders.single_file_model import FromOriginalModelMixin
 from diffusers.models.attention_processor import Attention
@@ -490,10 +492,17 @@ class BooguImageDoubleStreamTransformerBlock(nn.Module):
 
 
 class BooguImageTransformer2DModel(
-    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin
 ):
     """Boogu-Image transformer with mixed double-stream -> single-stream topology."""
 
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["double_stream_layers", "single_stream_layers"]
+
+
     _supports_gradient_checkpointing = True
     _no_split_modules = [
         "BooguImageTransformerBlock",

extensions_built_in/diffusion_models/chroma/chroma_model.py

@@ -5,16 +5,14 @@ import torch
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from PIL import Image
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
+from toolkit.models.v2.vae.autoencoder_kl import KLVAE
 from toolkit.basic import flush
-from diffusers import AutoencoderKL
 # from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
-from toolkit.dequantize import patch_dequantization_on_save
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze, QTensor
-from toolkit.util.quantize import quantize, get_qtype
-from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer
+from optimum.quanto import QTensor
 from .pipeline import ChromaPipeline, prepare_latent_image_ids
 from einops import rearrange, repeat
 import random
@@ -50,10 +48,12 @@ class FakeConfig:
         self.patch_size = 1
         
 class FakeCLIP(torch.nn.Module):
-    def __init__(self):
+    def __init__(self, device='cuda'):
         super().__init__()
         self.dtype = torch.bfloat16
-        self.device = 'cuda'
+        # the pipeline derives its execution device from this attribute;
+        # nn.Module.to() does not update it
+        self.device = device
         self.text_model = None
         self.tokenizer = None
         self.model_max_length = 77
@@ -65,6 +65,9 @@ class FakeCLIP(torch.nn.Module):
 class ChromaModel(BaseModel):
     arch = "chroma"
 
+    def get_transformer_block_names(self):
+        return ["double_blocks", "single_blocks"]
+
     def __init__(
             self,
             device,
@@ -148,83 +151,37 @@ class ChromaModel(BaseModel):
 
         self.print_and_status_update("Loading transformer")
         
-        chroma_state_dict = load_file(model_path, 'cpu')
-        
-        # determine number of double and single blocks
-        double_blocks = 0
-        single_blocks = 0
-        for key in chroma_state_dict.keys():
-            if "double_blocks" in key:
-                block_num = int(key.split(".")[1]) + 1
-                if block_num > double_blocks:
-                    double_blocks = block_num
-            elif "single_blocks" in key:
-                block_num = int(key.split(".")[1]) + 1
-                if block_num > single_blocks:
-                    single_blocks = block_num
-        print(f"Double Blocks: {double_blocks}")
-        print(f"Single Blocks: {single_blocks}")
-
-        chroma_params.depth = double_blocks
-        chroma_params.depth_single_blocks = single_blocks
-        transformer = Chroma(chroma_params)
-        
+        if model_path.endswith(".safetensors"):
+            transformer = Chroma.load_model(model_path, dtype=dtype)
+        else:
+            transformer = Chroma.load_from_state_dict(load_file(model_path, "cpu"), dtype)
         # add dtype, not sure why it doesnt have it
         transformer.dtype = dtype
-        # load the state dict into the model
-        transformer.load_state_dict(chroma_state_dict)
-        
-        transformer.to(self.quantize_device, dtype=dtype)
-        
-        transformer.config = FakeConfig()
-        transformer.config.num_layers = double_blocks
-        transformer.config.num_single_layers = single_blocks
 
-        if self.model_config.quantize:
-            # patch the state dict method
-            patch_dequantization_on_save(transformer)
-            quantization_type = get_qtype(self.model_config.qtype)
-            self.print_and_status_update("Quantizing transformer")
-            quantize(transformer, weights=quantization_type,
-                     **self.model_config.quantize_kwargs)
-            freeze(transformer)
-            transformer.to(self.device_torch)
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
+        transformer.config = FakeConfig()
+        transformer.config.num_layers = transformer.params.depth
+        transformer.config.num_single_layers = transformer.params.depth_single_blocks
+
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
 
         flush()
 
         self.print_and_status_update("Loading T5")
-        tokenizer_2 = T5TokenizerFast.from_pretrained(
-            extras_path, subfolder="tokenizer_2", torch_dtype=dtype
+        tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
+        text_encoder_2 = T5TextEncoder.load(
+            extras_path, **self.component_load_kwargs("te")
         )
-        text_encoder_2 = T5EncoderModel.from_pretrained(
-            extras_path, subfolder="text_encoder_2", torch_dtype=dtype
-        )
-        text_encoder_2.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing T5")
-            quantize(text_encoder_2, weights=get_qtype(
-                self.model_config.qtype))
-            freeze(text_encoder_2)
-            flush()
 
         # self.print_and_status_update("Loading CLIP")
-        text_encoder = FakeCLIP()
-        tokenizer = FakeCLIP()
+        text_encoder = FakeCLIP(device=self.device_torch)
+        tokenizer = FakeCLIP(device=self.device_torch)
         text_encoder.to(self.device_torch, dtype=dtype)
 
         self.noise_scheduler = ChromaModel.get_train_scheduler()
         
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKL.from_pretrained(
-            extras_path,
-            subfolder="vae",
-            torch_dtype=dtype
-        )
-        vae = vae.to(self.device_torch, dtype=dtype)
+        vae = KLVAE.load_model(extras_path, dtype=dtype, device=self.device_torch)
 
         self.print_and_status_update("Making pipe")
 
@@ -249,11 +206,13 @@ class ChromaModel(BaseModel):
         pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        # just to make sure everything is on the right device and dtype
-        text_encoder[0].to(self.device_torch)
+        # low_vram: text encoders stay on cpu; get_prompt_embeds moves them
+        # to the gpu on demand
+        if not self.low_vram:
+            text_encoder[0].to(self.device_torch)
+            text_encoder[1].to(self.device_torch)
         text_encoder[0].requires_grad_(False)
         text_encoder[0].eval()
-        text_encoder[1].to(self.device_torch)
         text_encoder[1].requires_grad_(False)
         text_encoder[1].eval()
         pipe.transformer = pipe.transformer.to(self.device_torch)
@@ -424,40 +383,24 @@ class ChromaModel(BaseModel):
         return self.text_encoder[1].encoder.block[0].layer[0].SelfAttention.q.weight.requires_grad
     
     def save_model(self, output_path, meta, save_dtype):
+        # comfy-format single-file save via the mixin (chroma's class keys ARE
+        # the original layout); handles torchao/Ostris dequant, not just quanto
         if not output_path.endswith(".safetensors"):
-            output_path =  output_path + ".safetensors"
-        # only save the unet
+            output_path = output_path + ".safetensors"
         transformer: Chroma = unwrap_model(self.model)
-        state_dict = transformer.state_dict()
-        save_dict = {}
-        for k, v in state_dict.items():
-            if isinstance(v, QTensor):
-                v = v.dequantize()
-            save_dict[k] = v.clone().to('cpu', dtype=save_dtype)
-        
-        meta = get_meta_for_safetensors(meta, name='chroma')
-        save_file(save_dict, output_path, metadata=meta)
+        transformer.save_model(
+            output_path,
+            dtype=save_dtype,
+            metadata=get_meta_for_safetensors(meta, name="chroma"),
+        )
 
     def get_loss_target(self, *args, **kwargs):
         noise = kwargs.get('noise')
         batch = kwargs.get('batch')
         return (noise - batch.latents).detach()
     
-    def convert_lora_weights_before_save(self, state_dict):
-        # currently starte with transformer. but needs to start with diffusion_model. for comfyui
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        # saved as diffusion_model. but needs to be transformer. for ai-toolkit
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
     
     def get_base_model_version(self):
         return "chroma"

extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py

@@ -5,16 +5,13 @@ import torch
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from PIL import Image
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
 from toolkit.basic import flush
-from diffusers import AutoencoderKL
 # from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
-from toolkit.dequantize import patch_dequantization_on_save
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze, QTensor
-from toolkit.util.quantize import quantize, get_qtype
-from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer
+from optimum.quanto import QTensor
 from .pipeline import ChromaPipeline, prepare_latent_image_ids
 from einops import rearrange, repeat
 import random
@@ -38,34 +35,16 @@ scheduler_config = {
     "use_dynamic_shifting": True
 }
 
-class FakeConfig:
-    # for diffusers compatability
-    def __init__(self):
-        self.attention_head_dim = 128
-        self.guidance_embeds = True
-        self.in_channels = 64
-        self.joint_attention_dim = 4096
-        self.num_attention_heads = 24
-        self.num_layers = 19
-        self.num_single_layers = 38
-        self.patch_size = 1
-        
-class FakeCLIP(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.dtype = torch.bfloat16
-        self.device = 'cuda'
-        self.text_model = None
-        self.tokenizer = None
-        self.model_max_length = 77
-
-    def forward(self, *args, **kwargs):
-        return torch.zeros(1, 1, 1).to(self.device)
+# shared with the base chroma model (identical stubs)
+from .chroma_model import FakeCLIP, FakeConfig
 
 
 class ChromaRadianceModel(BaseModel):
     arch = "chroma_radiance"
 
+    def get_transformer_block_names(self):
+        return ["double_blocks", "single_blocks"]
+
     def __init__(
             self,
             device,
@@ -152,73 +131,30 @@ class ChromaRadianceModel(BaseModel):
         
         if model_path.endswith('.pth') or model_path.endswith('.pt'):
             chroma_state_dict = torch.load(model_path, map_location='cpu', weights_only=True)
+            transformer = Chroma.load_from_state_dict(chroma_state_dict, dtype)
         else:
-            chroma_state_dict = load_file(model_path, 'cpu')
-        
-        # determine number of double and single blocks
-        double_blocks = 0
-        single_blocks = 0
-        for key in chroma_state_dict.keys():
-            if "double_blocks" in key:
-                block_num = int(key.split(".")[1]) + 1
-                if block_num > double_blocks:
-                    double_blocks = block_num
-            elif "single_blocks" in key:
-                block_num = int(key.split(".")[1]) + 1
-                if block_num > single_blocks:
-                    single_blocks = block_num
-        print(f"Double Blocks: {double_blocks}")
-        print(f"Single Blocks: {single_blocks}")
-
-        chroma_params.depth = double_blocks
-        chroma_params.depth_single_blocks = single_blocks
-        transformer = Chroma(chroma_params)
-        
+            transformer = Chroma.load_model(model_path, dtype=dtype)
         # add dtype, not sure why it doesnt have it
         transformer.dtype = dtype
-        # load the state dict into the model
-        transformer.load_state_dict(chroma_state_dict)
-        
-        transformer.to(self.quantize_device, dtype=dtype)
-        
-        transformer.config = FakeConfig()
-        transformer.config.num_layers = double_blocks
-        transformer.config.num_single_layers = single_blocks
 
-        if self.model_config.quantize:
-            # patch the state dict method
-            patch_dequantization_on_save(transformer)
-            quantization_type = get_qtype(self.model_config.qtype)
-            self.print_and_status_update("Quantizing transformer")
-            quantize(transformer, weights=quantization_type,
-                     **self.model_config.quantize_kwargs)
-            freeze(transformer)
-            transformer.to(self.device_torch)
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
+        transformer.config = FakeConfig()
+        transformer.config.num_layers = transformer.params.depth
+        transformer.config.num_single_layers = transformer.params.depth_single_blocks
+
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
 
         flush()
 
         self.print_and_status_update("Loading T5")
-        tokenizer_2 = T5TokenizerFast.from_pretrained(
-            extras_path, subfolder="tokenizer_2", torch_dtype=dtype
+        tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
+        text_encoder_2 = T5TextEncoder.load(
+            extras_path, **self.component_load_kwargs("te")
         )
-        text_encoder_2 = T5EncoderModel.from_pretrained(
-            extras_path, subfolder="text_encoder_2", torch_dtype=dtype
-        )
-        text_encoder_2.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing T5")
-            quantize(text_encoder_2, weights=get_qtype(
-                self.model_config.qtype))
-            freeze(text_encoder_2)
-            flush()
 
         # self.print_and_status_update("Loading CLIP")
-        text_encoder = FakeCLIP()
-        tokenizer = FakeCLIP()
+        text_encoder = FakeCLIP(device=self.device_torch)
+        tokenizer = FakeCLIP(device=self.device_torch)
         text_encoder.to(self.device_torch, dtype=dtype)
 
         self.noise_scheduler = ChromaRadianceModel.get_train_scheduler()
@@ -256,11 +192,13 @@ class ChromaRadianceModel(BaseModel):
         pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        # just to make sure everything is on the right device and dtype
-        text_encoder[0].to(self.device_torch)
+        # low_vram: text encoders stay on cpu; get_prompt_embeds moves them
+        # to the gpu on demand
+        if not self.low_vram:
+            text_encoder[0].to(self.device_torch)
+            text_encoder[1].to(self.device_torch)
         text_encoder[0].requires_grad_(False)
         text_encoder[0].eval()
-        text_encoder[1].to(self.device_torch)
         text_encoder[1].requires_grad_(False)
         text_encoder[1].eval()
         pipe.transformer = pipe.transformer.to(self.device_torch)
@@ -406,40 +344,24 @@ class ChromaRadianceModel(BaseModel):
         return False
     
     def save_model(self, output_path, meta, save_dtype):
+        # comfy-format single-file save via the mixin (chroma's class keys ARE
+        # the original layout); handles torchao/Ostris dequant, not just quanto
         if not output_path.endswith(".safetensors"):
-            output_path =  output_path + ".safetensors"
-        # only save the unet
+            output_path = output_path + ".safetensors"
         transformer: Chroma = unwrap_model(self.model)
-        state_dict = transformer.state_dict()
-        save_dict = {}
-        for k, v in state_dict.items():
-            if isinstance(v, QTensor):
-                v = v.dequantize()
-            save_dict[k] = v.clone().to('cpu', dtype=save_dtype)
-        
-        meta = get_meta_for_safetensors(meta, name='chroma')
-        save_file(save_dict, output_path, metadata=meta)
+        transformer.save_model(
+            output_path,
+            dtype=save_dtype,
+            metadata=get_meta_for_safetensors(meta, name="chroma"),
+        )
 
     def get_loss_target(self, *args, **kwargs):
         noise = kwargs.get('noise')
         batch = kwargs.get('batch')
         return (noise - batch.latents).detach()
     
-    def convert_lora_weights_before_save(self, state_dict):
-        # currently starte with transformer. but needs to start with diffusion_model. for comfyui
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        # saved as diffusion_model. but needs to be transformer. for ai-toolkit
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
     
     def get_base_model_version(self):
         return "chroma_radiance"

extensions_built_in/diffusion_models/chroma/src/model.py

@@ -1,4 +1,6 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 
 import torch
 from torch import Tensor, nn
@@ -86,11 +88,40 @@ def modify_mask_to_attend_padding(mask, max_seq_length, num_extra_padding=8):
     return modified_mask
 
 
-class Chroma(nn.Module):
+class Chroma(nn.Module, OstrisModelMixin):
     """
     Transformer model for flow matching on sequences.
     """
 
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict):
+        # block counts come from the checkpoint's key indices
+        double_blocks = 0
+        single_blocks = 0
+        for key in state_dict.keys():
+            if "double_blocks" in key:
+                block_num = int(key.split(".")[1]) + 1
+                if block_num > double_blocks:
+                    double_blocks = block_num
+            elif "single_blocks" in key:
+                block_num = int(key.split(".")[1]) + 1
+                if block_num > single_blocks:
+                    single_blocks = block_num
+        print(f"Double Blocks: {double_blocks}")
+        print(f"Single Blocks: {single_blocks}")
+        return replace(
+            chroma_params, depth=double_blocks, depth_single_blocks=single_blocks
+        )
+
+    @classmethod
+    def aitk_from_config(cls, config):
+        with torch.device("meta"):
+            return cls(config)
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["double_blocks", "single_blocks"]
+
     def __init__(self, params: ChromaParams):
         super().__init__()
         self.params = params

extensions_built_in/diffusion_models/chroma/src/radiance.py

@@ -1,4 +1,6 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 
 import torch
 from torch import Tensor, nn
@@ -100,11 +102,40 @@ def modify_mask_to_attend_padding(mask, max_seq_length, num_extra_padding=8):
     return modified_mask
 
 
-class Chroma(nn.Module):
+class Chroma(nn.Module, OstrisModelMixin):
     """
     Transformer model for flow matching on sequences.
     """
 
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict):
+        # block counts come from the checkpoint's key indices
+        double_blocks = 0
+        single_blocks = 0
+        for key in state_dict.keys():
+            if "double_blocks" in key:
+                block_num = int(key.split(".")[1]) + 1
+                if block_num > double_blocks:
+                    double_blocks = block_num
+            elif "single_blocks" in key:
+                block_num = int(key.split(".")[1]) + 1
+                if block_num > single_blocks:
+                    single_blocks = block_num
+        print(f"Double Blocks: {double_blocks}")
+        print(f"Single Blocks: {single_blocks}")
+        return replace(
+            chroma_params, depth=double_blocks, depth_single_blocks=single_blocks
+        )
+
+    @classmethod
+    def aitk_from_config(cls, config):
+        with torch.device("meta"):
+            return cls(config)
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["double_blocks", "single_blocks"]
+
     def __init__(self, params: ChromaParams):
         super().__init__()
         self.params = params

extensions_built_in/diffusion_models/ernie_image/ernie_image.py

@@ -5,15 +5,14 @@ import torch
 import yaml
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.mistral3 import Mistral3ModelEncoder
+from toolkit.models.v2.vae.autoencoder_kl_flux2 import Flux2KLVAE
 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 optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
-from toolkit.memory_management import MemoryManager
 
 from transformers import AutoTokenizer, AutoModel
 
@@ -79,75 +78,31 @@ class ErnieImageModel(BaseModel):
 
         self.print_and_status_update("Loading transformer")
 
-        transformer_path = model_path
-        transformer_subfolder = "transformer"
-        if os.path.exists(transformer_path):
-            transformer_subfolder = None
-            transformer_path = os.path.join(transformer_path, "transformer")
+        if os.path.exists(model_path):
             # check if the path is a full checkpoint.
             te_folder_path = os.path.join(model_path, "text_encoder")
             # if we have the te, this folder is a full checkpoint, use it as the base
             if os.path.exists(te_folder_path):
                 base_model_path = model_path
 
-        transformer = ErnieImageTransformer2DModel.from_pretrained(
-            transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype
+        # load + quantize + offload + placement, all driven by model_config
+        transformer = ErnieImageTransformer2DModel.load(
+            model_path, **self.component_load_kwargs("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,
-                ignore_modules=[
-                    transformer.x_embedder,
-                ],
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
-
         flush()
 
         self.print_and_status_update("Text Encoder")
         tokenizer = AutoTokenizer.from_pretrained(
             base_model_path, subfolder="tokenizer", torch_dtype=dtype
         )
-        text_encoder = AutoModel.from_pretrained(
-            base_model_path, subfolder="text_encoder", torch_dtype=dtype
+        text_encoder = Mistral3ModelEncoder.load(
+            base_model_path, subfolder="text_encoder", **self.component_load_kwargs("te")
         )
-
-        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,
-            )
-
-        text_encoder.to(self.device_torch, dtype=dtype)
         flush()
 
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
-
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKLFlux2.from_pretrained(
-            base_model_path, subfolder="vae", torch_dtype=dtype
+        vae = Flux2KLVAE.load_model(            base_model_path, dtype=dtype
         ).to(self.device_torch, dtype=dtype)
 
         self.noise_scheduler = ErnieImageModel.get_train_scheduler()
@@ -178,8 +133,10 @@ class ErnieImageModel(BaseModel):
             pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        # just to make sure everything is on the right device and dtype
-        text_encoder[0].to(self.device_torch)
+        # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
+        # to the gpu on demand
+        if not self.low_vram:
+            text_encoder[0].to(self.device_torch)
         text_encoder[0].requires_grad_(False)
         text_encoder[0].eval()
         flush()
@@ -377,16 +334,5 @@ class ErnieImageModel(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["layers"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd

extensions_built_in/diffusion_models/ernie_image/transformer.py

@@ -33,6 +33,8 @@ from diffusers.models.attention_dispatch import dispatch_attention_fn
 from diffusers.models.attention_processor import Attention
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_utils import ModelMixin
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 from diffusers.models.normalization import RMSNorm
 
 
@@ -288,8 +290,17 @@ class ErnieImageAdaLNContinuous(nn.Module):
         return x
 
 
-class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin):
+class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin, OstrisModelMixin):
     _supports_gradient_checkpointing = True
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]
+
+    def get_offload_ignore_modules(self):
+        return [self.x_embedder]
+
 
     @register_to_config
     def __init__(

extensions_built_in/diffusion_models/example_model/example_model.py

@@ -32,7 +32,7 @@ from toolkit.models.base_model import BaseModel
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
+from toolkit.util.quantize import quantize, get_qtype
 
 from .src.model import ExampleTransformer2DModel
 from .src.pipeline import ExamplePipeline, pad_prompt_embeds
@@ -165,21 +165,14 @@ class ExampleModel(BaseModel):
         del state_dict
         flush()  # gc + empty cuda cache; call it after dropping anything big
 
-        if self.model_config.quantize:
-            # quantize_model handles qtype selection, exclusions and device
-            # juggling, and leaves the model on CPU
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            flush()
-
-        if self.model_config.low_vram:
-            # leave it on CPU; get_noise_prediction moves it over when needed
-            transformer.to("cpu")
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
+        # quantize + offload + placement, all driven by model_config:
+        # component_load_kwargs derives qtype (incl. an accuracy recovery
+        # adapter), the layer-offload fraction and the target device
+        # (low_vram parks on CPU); aitk_post_load applies them. Models whose
+        # checkpoint sourcing is standard can collapse the build + this into
+        # one call: ExampleTransformer2DModel.load(path, **kwargs).
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
-        # For partial layer offloading support see MemoryManager.attach usage
-        # in ../ideogram4/ideogram4.py or ../z_image/z_image.py.
 
         # --- text encoder + tokenizer (stock transformers model) ---
         self.print_and_status_update("Loading text encoder")
@@ -504,19 +497,8 @@ class ExampleModel(BaseModel):
         attribute in src/model.py."""
         return ["blocks"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        """Map internal LoRA keys to the ecosystem-standard naming right before
-        the .safetensors is written. Most modern models ship LoRAs with a
-        ``diffusion_model.`` prefix (ComfyUI convention); internally ai-toolkit
-        uses ``transformer.``."""
-        return {
-            k.replace("transformer.", "diffusion_model."): v
-            for k, v in state_dict.items()
-        }
+    # LoRA keys save with the ecosystem-standard ``diffusion_model.`` prefix
+    # (ComfyUI convention) and load back to the internal ``transformer.``
+    # prefix; see BaseModel.convert_lora_weights_before_save/load
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        """Inverse of the above, applied when resuming from a saved LoRA."""
-        return {
-            k.replace("diffusion_model.", "transformer."): v
-            for k, v in state_dict.items()
-        }

extensions_built_in/diffusion_models/example_model/src/model.py

@@ -51,6 +51,8 @@ import torch.nn.functional as F
 from torch import nn
 from torch.utils.checkpoint import checkpoint
 
+from toolkit.models.v2._mixin import OstrisModelMixin
+
 
 def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
     """Standard sinusoidal embedding.
@@ -135,8 +137,16 @@ class ExampleTransformerBlock(nn.Module):
         return x
 
 
-class ExampleTransformer2DModel(nn.Module):
-    """The denoiser. Plain ``nn.Module`` on purpose.
+class ExampleTransformer2DModel(nn.Module, OstrisModelMixin):
+    """The denoiser. Plain ``nn.Module`` plus ``OstrisModelMixin``.
+
+    The mixin supplies the universal component API every v2 module shares:
+    ``load()`` / ``load_model()`` / ``aitk_post_load()`` (quantize, layer
+    offloading and device placement driven by the holder's model_config via
+    ``BaseModel.component_load_kwargs(role)``), plus comfy-format save/load.
+    Override its class hooks (``get_transformer_block_names``,
+    ``get_quantization_exclude_modules``, ``get_offload_ignore_modules``,
+    ``convert_state_dict_on_load/save``) as needed.
 
     You could also subclass ``diffusers.ModelMixin``/``ConfigMixin`` (see
     ../../ernie_image/transformer.py) to get ``save_pretrained``,
@@ -153,6 +163,12 @@ class ExampleTransformer2DModel(nn.Module):
     names when deciding where to attach LoRA layers.
     """
 
+    @classmethod
+    def get_transformer_block_names(cls):
+        # attribute name(s) of the repeated-block ModuleList(s); the quantizer
+        # streams these blocks through the GPU one at a time
+        return ["blocks"]
+
     def __init__(
         self,
         in_channels: int = 16,    # VAE latent channels
@@ -183,8 +199,8 @@ class ExampleTransformer2DModel(nn.Module):
         )
 
         # ``blocks`` is the repeated-layer ModuleList. The attribute name is
-        # what ExampleModel.get_transformer_block_names() returns, which the
-        # LoRA code uses for block targeting / "transformer only" training.
+        # what get_transformer_block_names() returns, which the LoRA code uses
+        # for block targeting and the quantizer uses for block streaming.
         self.blocks = nn.ModuleList(
             [
                 ExampleTransformerBlock(hidden_size, num_heads)

extensions_built_in/diffusion_models/f_light/f_light.py

@@ -6,15 +6,14 @@ import yaml
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from PIL import Image
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
+from toolkit.models.v2.vae.autoencoder_kl import KLVAE
 from toolkit.basic import flush
-from diffusers import AutoencoderKL
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
-from toolkit.dequantize import patch_dequantization_on_save
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze, QTensor
-from toolkit.util.quantize import quantize, get_qtype
-from transformers import T5TokenizerFast, T5EncoderModel
+from optimum.quanto import QTensor
+
 from .src import FLitePipeline, DiT
 
 if TYPE_CHECKING:
@@ -34,6 +33,9 @@ scheduler_config = {
 class FLiteModel(BaseModel):
     arch = "f-lite"
 
+    def get_transformer_block_names(self):
+        return ["blocks"]
+
     def __init__(
             self,
             device,
@@ -74,54 +76,22 @@ class FLiteModel(BaseModel):
 
         self.print_and_status_update("Loading transformer")
 
-        transformer = DiT.from_pretrained(
-            model_path,
-            subfolder="dit_model",
-            torch_dtype=dtype,
+        transformer = DiT.load(
+            model_path, **self.component_load_kwargs("transformer")
         )
-        
-        transformer.to(self.quantize_device, dtype=dtype)
-
-        if self.model_config.quantize:
-            # patch the state dict method
-            patch_dequantization_on_save(transformer)
-            quantization_type = get_qtype(self.model_config.qtype)
-            self.print_and_status_update("Quantizing transformer")
-            quantize(transformer, weights=quantization_type,
-                     **self.model_config.quantize_kwargs)
-            freeze(transformer)
-            transformer.to(self.device_torch)
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
 
         flush()
 
         self.print_and_status_update("Loading T5")
-        tokenizer = T5TokenizerFast.from_pretrained(
-            extras_path, subfolder="tokenizer", torch_dtype=dtype
+        tokenizer = T5TextEncoder.load_tokenizer(extras_path, subfolder="tokenizer")
+        text_encoder = T5TextEncoder.load(
+            extras_path, subfolder="text_encoder", **self.component_load_kwargs("te")
         )
-        text_encoder = T5EncoderModel.from_pretrained(
-            extras_path, subfolder="text_encoder", torch_dtype=dtype
-        )
-        text_encoder.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing T5")
-            quantize(text_encoder, weights=get_qtype(
-                self.model_config.qtype))
-            freeze(text_encoder)
-            flush()
 
         self.noise_scheduler = FLiteModel.get_train_scheduler()
         
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKL.from_pretrained(
-            extras_path,
-            subfolder="vae",
-            torch_dtype=dtype
-        )
-        vae = vae.to(self.device_torch, dtype=dtype)
+        vae = KLVAE.load_model(extras_path, dtype=dtype, device=self.device_torch)
 
         self.print_and_status_update("Making pipe")
 
@@ -145,8 +115,10 @@ class FLiteModel(BaseModel):
         pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        # just to make sure everything is on the right device and dtype
-        text_encoder[0].to(self.device_torch)
+        # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
+        # to the gpu on demand
+        if not self.low_vram:
+            text_encoder[0].to(self.device_torch)
         text_encoder[0].requires_grad_(False)
         text_encoder[0].eval()
         pipe.transformer = pipe.transformer.to(self.device_torch)
@@ -270,21 +242,8 @@ class FLiteModel(BaseModel):
         # return (noise - batch.latents).detach()
         return (batch.latents - noise).detach()
     
-    def convert_lora_weights_before_save(self, state_dict):
-        # currently starte with transformer. but needs to start with diffusion_model. for comfyui
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        # saved as diffusion_model. but needs to be transformer. for ai-toolkit
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
     
     def get_base_model_version(self):
         return "f-lite"

extensions_built_in/diffusion_models/f_light/src/model.py

@@ -7,6 +7,8 @@ import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
 from diffusers.models.modeling_utils import ModelMixin
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 from diffusers.utils.accelerate_utils import apply_forward_hook
 from einops import rearrange
 from peft import get_peft_model_state_dict, set_peft_model_state_dict
@@ -302,7 +304,13 @@ def apply_rotary_emb(x, cos, sin):
     return torch.cat([y1, y2], 3).to(dtype=orig_dtype)
 
 
-class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin):  # type: ignore[misc]
+class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin, OstrisModelMixin):  # type: ignore[misc]
+    aitk_subfolder = "dit_model"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["blocks"]
+
     _supports_gradient_checkpointing = True
     
     @register_to_config

extensions_built_in/diffusion_models/flux2/flux2_klein_model.py

@@ -1,9 +1,7 @@
 from .flux2_model import Flux2Model
 from transformers import Qwen3ForCausalLM, Qwen2Tokenizer
-from optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype
+from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
 from toolkit.config_modules import ModelConfig
-from toolkit.memory_management.manager import MemoryManager
 from toolkit.basic import flush
 from .src.model import Klein9BParams, Klein4BParams
 
@@ -40,28 +38,11 @@ class Flux2KleinModel(Flux2Model):
         dtype = self.torch_dtype
         self.print_and_status_update("Loading Qwen3")
 
-        text_encoder: Qwen3ForCausalLM = Qwen3ForCausalLM.from_pretrained(
-            self.flux2_klein_te_path,
-            torch_dtype=dtype,
+        # load + quantize + offload + placement, all driven by model_config
+        text_encoder = Qwen3TextEncoder.load(
+            self.flux2_klein_te_path, subfolder="", **self.component_load_kwargs("te")
         )
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Qwen3")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
-        elif not self.model_config.low_vram:
-            text_encoder.to(self.device_torch, dtype=dtype)
-            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,
-            )
+        flush()
 
         tokenizer = Qwen2Tokenizer.from_pretrained(self.flux2_klein_te_path)
         return text_encoder, tokenizer

extensions_built_in/diffusion_models/flux2/flux2_model.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, List, Optional
 import huggingface_hub
 import torch
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
-from toolkit.memory_management.manager import MemoryManager
 from toolkit.metadata import get_meta_for_safetensors
 from toolkit.models.base_model import BaseModel
 from toolkit.basic import flush
@@ -13,15 +12,18 @@ from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
-from toolkit.dequantize import patch_dequantization_on_save
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze, QTensor
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
+from optimum.quanto import QTensor
 
 from transformers import AutoProcessor, Mistral3ForConditionalGeneration
+from toolkit.models.v2.text_encoders.mistral3 import Mistral3TextEncoder
 from .src.model import Flux2, Flux2Params
 from .src.pipeline import Flux2Pipeline
-from .src.autoencoder import AutoEncoder, AutoEncoderParams, AutoEncoderSmallDecoderParams
+from toolkit.models.v2.vae.flux2_kl import (
+    AutoEncoder,
+    AutoEncoderParams,
+    AutoEncoderSmallDecoderParams,
+)
 from safetensors.torch import load_file, save_file
 from PIL import Image
 import torch.nn.functional as F
@@ -95,33 +97,22 @@ class Flux2Model(BaseModel):
         dtype = self.torch_dtype
         self.print_and_status_update("Loading Mistral")
 
-        text_encoder: Mistral3ForConditionalGeneration = (
-            Mistral3ForConditionalGeneration.from_pretrained(
-                MISTRAL_PATH,
-                torch_dtype=dtype,
-            )
+        # load + quantize + offload + placement, all driven by model_config
+        # tie_word_embeddings=False: the checkpoint carries both embed_tokens
+        # and lm_head with different values; the config's tie claim is wrong
+        text_encoder = Mistral3TextEncoder.load(
+            MISTRAL_PATH,
+            subfolder="",
+            tie_word_embeddings=False,
+            **self.component_load_kwargs("te"),
         )
-        text_encoder.to(self.device_torch, dtype=dtype)
-
         flush()
 
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Mistral")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype))
-            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,
-            )
-
-        tokenizer = AutoProcessor.from_pretrained(MISTRAL_PATH)
+        # fix_mistral_regex=False: keep the exact tokenization flux2 has always
+        # used (True would change the pre-tokenizer and shift conditioning)
+        tokenizer = AutoProcessor.from_pretrained(
+            MISTRAL_PATH, fix_mistral_regex=False
+        )
         return text_encoder, tokenizer
 
     def load_model(self):
@@ -132,9 +123,6 @@ class Flux2Model(BaseModel):
         transformer_path = model_path
 
         self.print_and_status_update("Loading transformer")
-        with torch.device("meta"):
-            transformer = Flux2(self.get_flux2_params())
-
         # use local path if provided
         if os.path.exists(os.path.join(transformer_path, self.flux2_te_filename)):
             transformer_path = os.path.join(transformer_path, self.flux2_te_filename)
@@ -148,39 +136,14 @@ class Flux2Model(BaseModel):
             )
 
         transformer_state_dict = load_file(transformer_path, device="cpu")
+        transformer = Flux2.load_from_state_dict(
+            transformer_state_dict, dtype, config=self.get_flux2_params()
+        )
 
-        # cast to dtype
-        for key in transformer_state_dict:
-            transformer_state_dict[key] = transformer_state_dict[key].to(dtype)
-
-        transformer.load_state_dict(transformer_state_dict, assign=True)
-
-        if self.model_config.quantize:
-            # patch the state dict method
-            patch_dequantization_on_save(transformer)
-            # Avoid full-model peak VRAM allocation before quantization.
-            self.print_and_status_update("Keeping transformer on CPU for quantization")
-            self.print_and_status_update("Quantizing Transformer")
-            quantize_model(self, transformer)
-            flush()
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("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")
-
         text_encoder, tokenizer = self.load_te()
 
         self.print_and_status_update("Loading VAE")
@@ -207,21 +170,8 @@ class Flux2Model(BaseModel):
                 token=HF_TOKEN,
             )
         
-        vae_state_dict = load_file(vae_path, device="cpu")
-        
-        autoencoder_params = AutoEncoderParams()
-        if vae_state_dict['decoder.up.0.block.0.conv1.bias'].shape[0] == 96:
-            # this is the small decoder version
-            autoencoder_params = AutoEncoderSmallDecoderParams()
-        
-        with torch.device("meta"):
-            vae = AutoEncoder(autoencoder_params)
-
-        # cast to dtype
-        for key in vae_state_dict:
-            vae_state_dict[key] = vae_state_dict[key].to(dtype)
-
-        vae.load_state_dict(vae_state_dict, assign=True)
+        # config sniffed from the checkpoint (small-decoder detection)
+        vae = AutoEncoder.load_model(vae_path, dtype=dtype)
 
         self.noise_scheduler = Flux2Model.get_train_scheduler()
 
@@ -505,19 +455,7 @@ class Flux2Model(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["double_blocks", "single_blocks"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
-
-    def convert_lora_weights_before_load(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
     def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
         if device is None:

extensions_built_in/diffusion_models/flux2/src/autoencoder.py

@@ -1,435 +0,0 @@
-from dataclasses import dataclass, field
-
-import torch
-from einops import rearrange
-from torch import Tensor, nn
-import math
-import torch.utils.checkpoint as ckpt
-
-
-@dataclass
-class AutoEncoderParams:
-    resolution: int = 256
-    in_channels: int = 3
-    ch: int = 128
-    out_ch: int = 3
-    ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
-    num_res_blocks: int = 2
-    z_channels: int = 32
-
-@dataclass
-class AutoEncoderSmallDecoderParams:
-    resolution: int = 256
-    in_channels: int = 3
-    ch: int = 128
-    ch_encoder: int = 96
-    out_ch: int = 3
-    ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
-    num_res_blocks: int = 2
-    z_channels: int = 32
-
-
-def swish(x: Tensor) -> Tensor:
-    return x * torch.sigmoid(x)
-
-
-class AttnBlock(nn.Module):
-    def __init__(self, in_channels: int):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = nn.GroupNorm(
-            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
-        )
-
-        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-
-    def attention(self, h_: Tensor) -> Tensor:
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        b, c, h, w = q.shape
-        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
-        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
-        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
-        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
-
-        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return x + self.proj_out(self.attention(x))
-
-
-class ResnetBlock(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-
-        self.norm1 = nn.GroupNorm(
-            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
-        )
-        self.conv1 = nn.Conv2d(
-            in_channels, out_channels, kernel_size=3, stride=1, padding=1
-        )
-        self.norm2 = nn.GroupNorm(
-            num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
-        )
-        self.conv2 = nn.Conv2d(
-            out_channels, out_channels, kernel_size=3, stride=1, padding=1
-        )
-        if self.in_channels != self.out_channels:
-            self.nin_shortcut = nn.Conv2d(
-                in_channels, out_channels, kernel_size=1, stride=1, padding=0
-            )
-
-    def forward(self, x):
-        h = x
-        h = self.norm1(h)
-        h = swish(h)
-        h = self.conv1(h)
-
-        h = self.norm2(h)
-        h = swish(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            x = self.nin_shortcut(x)
-
-        return x + h
-
-
-class Downsample(nn.Module):
-    def __init__(self, in_channels: int):
-        super().__init__()
-        # no asymmetric padding in torch conv, must do it ourselves
-        self.conv = nn.Conv2d(
-            in_channels, in_channels, kernel_size=3, stride=2, padding=0
-        )
-
-    def forward(self, x: Tensor):
-        pad = (0, 1, 0, 1)
-        x = nn.functional.pad(x, pad, mode="constant", value=0)
-        x = self.conv(x)
-        return x
-
-
-class Upsample(nn.Module):
-    def __init__(self, in_channels: int):
-        super().__init__()
-        self.conv = nn.Conv2d(
-            in_channels, in_channels, kernel_size=3, stride=1, padding=1
-        )
-
-    def forward(self, x: Tensor):
-        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        x = self.conv(x)
-        return x
-
-
-class Encoder(nn.Module):
-    def __init__(
-        self,
-        resolution: int,
-        in_channels: int,
-        ch: int,
-        ch_mult: list[int],
-        num_res_blocks: int,
-        z_channels: int,
-    ):
-        super().__init__()
-        self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
-        self.ch = ch
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        # downsampling
-        self.conv_in = nn.Conv2d(
-            in_channels, self.ch, kernel_size=3, stride=1, padding=1
-        )
-
-        curr_res = resolution
-        in_ch_mult = (1,) + tuple(ch_mult)
-        self.in_ch_mult = in_ch_mult
-        self.down = nn.ModuleList()
-        block_in = self.ch
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch * in_ch_mult[i_level]
-            block_out = ch * ch_mult[i_level]
-            for _ in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
-                block_in = block_out
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions - 1:
-                down.downsample = Downsample(block_in)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-
-        # end
-        self.norm_out = nn.GroupNorm(
-            num_groups=32, num_channels=block_in, eps=1e-6, affine=True
-        )
-        self.conv_out = nn.Conv2d(
-            block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1
-        )
-        self.gradient_checkpointing = False
-
-    def enable_gradient_checkpointing(self):
-        self.gradient_checkpointing = True
-
-    def forward(self, x: Tensor) -> Tensor:
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                if torch.is_grad_enabled() and self.gradient_checkpointing:
-                    h = ckpt.checkpoint(self.down[i_level].block[i_block], hs[-1])
-                    if len(self.down[i_level].attn) > 0:
-                        h = ckpt.checkpoint(self.down[i_level].attn[i_block], h)
-                else:
-                    h = self.down[i_level].block[i_block](hs[-1])
-                    if len(self.down[i_level].attn) > 0:
-                        h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions - 1:
-                if torch.is_grad_enabled() and self.gradient_checkpointing:
-                    hs.append(ckpt.checkpoint(self.down[i_level].downsample, hs[-1]))
-                else:
-                    hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-            h = ckpt.checkpoint(self.mid.block_1, h)
-            h = ckpt.checkpoint(self.mid.attn_1, h)
-            h = ckpt.checkpoint(self.mid.block_2, h)
-        else:
-            h = self.mid.block_1(h)
-            h = self.mid.attn_1(h)
-            h = self.mid.block_2(h)
-        # end
-        h = self.norm_out(h)
-        h = swish(h)
-        h = self.conv_out(h)
-        h = self.quant_conv(h)
-        return h
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        ch: int,
-        out_ch: int,
-        ch_mult: list[int],
-        num_res_blocks: int,
-        in_channels: int,
-        resolution: int,
-        z_channels: int,
-    ):
-        super().__init__()
-        self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
-        self.ch = ch
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.ffactor = 2 ** (self.num_resolutions - 1)
-
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        block_in = ch * ch_mult[self.num_resolutions - 1]
-        curr_res = resolution // 2 ** (self.num_resolutions - 1)
-        self.z_shape = (1, z_channels, curr_res, curr_res)
-
-        # z to block_in
-        self.conv_in = nn.Conv2d(
-            z_channels, block_in, kernel_size=3, stride=1, padding=1
-        )
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch * ch_mult[i_level]
-            for _ in range(self.num_res_blocks + 1):
-                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
-                block_in = block_out
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in)
-                curr_res = curr_res * 2
-            self.up.insert(0, up)  # prepend to get consistent order
-
-        # end
-        self.norm_out = nn.GroupNorm(
-            num_groups=32, num_channels=block_in, eps=1e-6, affine=True
-        )
-        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
-        self.gradient_checkpointing = False
-
-    def enable_gradient_checkpointing(self):
-        self.gradient_checkpointing = True
-
-    def forward(self, z: Tensor) -> Tensor:
-        z = self.post_quant_conv(z)
-
-        # get dtype for proper tracing
-        upscale_dtype = next(self.up.parameters()).dtype
-
-        # z to block_in
-        h = self.conv_in(z)
-
-        # middle
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-            h = ckpt.checkpoint(self.mid.block_1, h)
-            h = ckpt.checkpoint(self.mid.attn_1, h)
-            h = ckpt.checkpoint(self.mid.block_2, h)
-        else:
-            h = self.mid.block_1(h)
-            h = self.mid.attn_1(h)
-            h = self.mid.block_2(h)
-
-        # cast to proper dtype
-        h = h.to(upscale_dtype)
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks + 1):
-                if torch.is_grad_enabled() and self.gradient_checkpointing:
-                    h = ckpt.checkpoint(self.up[i_level].block[i_block], h)
-                    if len(self.up[i_level].attn) > 0:
-                        h = ckpt.checkpoint(self.up[i_level].attn[i_block], h)
-                else:
-                    h = self.up[i_level].block[i_block](h)
-                    if len(self.up[i_level].attn) > 0:
-                        h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                if torch.is_grad_enabled() and self.gradient_checkpointing:
-                    h = ckpt.checkpoint(self.up[i_level].upsample, h)
-                else:
-                    h = self.up[i_level].upsample(h)
-
-        # end
-        h = self.norm_out(h)
-        h = swish(h)
-        h = self.conv_out(h)
-        return h
-
-
-class AutoEncoder(nn.Module):
-    def __init__(self, params: AutoEncoderParams):
-        super().__init__()
-        self.params = params
-        self.encoder = Encoder(
-            resolution=params.resolution,
-            in_channels=params.in_channels,
-            ch=params.ch,
-            ch_mult=params.ch_mult,
-            num_res_blocks=params.num_res_blocks,
-            z_channels=params.z_channels,
-        )
-        decoder_ch = params.ch
-        if hasattr(params, "ch_encoder"):
-            decoder_ch = params.ch_encoder
-        self.decoder = Decoder(
-            resolution=params.resolution,
-            in_channels=params.in_channels,
-            ch=decoder_ch,
-            out_ch=params.out_ch,
-            ch_mult=params.ch_mult,
-            num_res_blocks=params.num_res_blocks,
-            z_channels=params.z_channels,
-        )
-
-        self.bn_eps = 1e-4
-        self.bn_momentum = 0.1
-        self.ps = [2, 2]
-        self.bn = torch.nn.BatchNorm2d(
-            math.prod(self.ps) * params.z_channels,
-            eps=self.bn_eps,
-            momentum=self.bn_momentum,
-            affine=False,
-            track_running_stats=True,
-        )
-        self._gradient_checkpointing = False
-
-    @property
-    def gradient_checkpointing(self):
-        return self._gradient_checkpointing
-
-    @gradient_checkpointing.setter
-    def gradient_checkpointing(self, value: bool):
-        self._gradient_checkpointing = value
-        self.encoder.gradient_checkpointing = value
-        self.decoder.gradient_checkpointing = value
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
-    @property
-    def dtype(self):
-        return next(self.parameters()).dtype
-
-    def enable_gradient_checkpointing(self):
-        self.gradient_checkpointing = True
-        self.encoder.enable_gradient_checkpointing()
-        self.decoder.enable_gradient_checkpointing()
-
-    def normalize(self, z):
-        self.bn.eval()
-        return self.bn(z)
-
-    def inv_normalize(self, z):
-        self.bn.eval()
-        s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn_eps)
-        m = self.bn.running_mean.view(1, -1, 1, 1)
-        return z * s + m
-
-    def encode(self, x: Tensor) -> Tensor:
-        moments = self.encoder(x)
-        mean = torch.chunk(moments, 2, dim=1)[0]
-
-        z = rearrange(
-            mean,
-            "... c (i pi) (j pj)  -> ... (c pi pj) i j",
-            pi=self.ps[0],
-            pj=self.ps[1],
-        )
-        z = self.normalize(z)
-        return z
-
-    def decode(self, z: Tensor) -> Tensor:
-        z = self.inv_normalize(z)
-        z = rearrange(
-            z,
-            "... (c pi pj) i j -> ... c (i pi) (j pj)",
-            pi=self.ps[0],
-            pj=self.ps[1],
-        )
-        dec = self.decoder(z)
-        return dec

extensions_built_in/diffusion_models/flux2/src/model.py

@@ -1,4 +1,6 @@
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 from einops import rearrange
 from torch import Tensor, nn
 import torch.utils.checkpoint as ckpt
@@ -54,7 +56,11 @@ class FakeConfig:
         self.patch_size = 1
 
 
-class Flux2(nn.Module):
+class Flux2(nn.Module, OstrisModelMixin):
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["double_blocks", "single_blocks"]
+
     def __init__(self, params: Flux2Params):
         super().__init__()
         self.config = FakeConfig()

extensions_built_in/diffusion_models/flux2/src/pipeline.py

@@ -12,7 +12,7 @@ from diffusers.utils import (
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.utils import BaseOutput
-from .autoencoder import AutoEncoder
+from toolkit.models.v2.vae.flux2_kl import AutoEncoder
 from .model import Flux2
 from einops import rearrange
 from transformers import AutoProcessor, Mistral3ForConditionalGeneration
@@ -109,15 +109,19 @@ class Flux2Pipeline(DiffusionPipeline):
         # Process all messages at once
         # with image processing a too short max length can throw an error in here.
         try:
+            # tokenization kwargs ride in processor_kwargs (same values end up
+            # in the same place; loose **kwargs just warn on new transformers)
             inputs = self.tokenizer.apply_chat_template(
                 messages_batch,
                 add_generation_prompt=False,
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt",
-                padding="max_length",
-                truncation=True,
-                max_length=max_sequence_length,
+                processor_kwargs={
+                    "padding": "max_length",
+                    "truncation": True,
+                    "max_length": max_sequence_length,
+                },
             )
         except ValueError as e:
             print(

extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py

@@ -8,17 +8,19 @@ from toolkit import train_tools
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from PIL import Image
 from toolkit.models.base_model import BaseModel
-from diffusers import FluxTransformer2DModel, AutoencoderKL, FluxKontextPipeline
+from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
+from toolkit.models.v2.text_encoders.clip import CLIPTextEncoder
+from toolkit.models.v2.vae.autoencoder_kl import KLVAE
+from diffusers import FluxKontextPipeline
+from toolkit.models.v2.diffusion_models.flux import FluxTransformer2DModel
 from toolkit.basic import flush
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
 from toolkit.models.flux import add_model_gpu_splitter_to_flux, bypass_flux_guidance, restore_flux_guidance
-from toolkit.dequantize import patch_dequantization_on_save
 from toolkit.accelerator import get_accelerator, unwrap_model
-from optimum.quanto import freeze, QTensor
+from optimum.quanto import QTensor
 from toolkit.util.mask import generate_random_mask, random_dialate_mask
-from toolkit.util.quantize import quantize, get_qtype
-from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer
+
 from einops import rearrange, repeat
 import random
 import torch.nn.functional as F
@@ -36,11 +38,12 @@ scheduler_config = {
     "use_dynamic_shifting": True
 }
 
-
-
 class FluxKontextModel(BaseModel):
     arch = "flux_kontext"
 
+    def get_transformer_block_names(self):
+        return ["transformer_blocks", "single_transformer_blocks"]
+
     def __init__(
             self,
             device,
@@ -80,11 +83,7 @@ class FluxKontextModel(BaseModel):
         # so we need this for the VAE, te, etc
         base_model_path = self.model_config.extras_name_or_path
 
-        transformer_path = model_path
-        transformer_subfolder = 'transformer'
-        if os.path.exists(transformer_path):
-            transformer_subfolder = None
-            transformer_path = os.path.join(transformer_path, 'transformer')
+        if os.path.exists(model_path):
             # check if the path is a full checkpoint.
             te_folder_path = os.path.join(model_path, 'text_encoder')
             # if we have the te, this folder is a full checkpoint, use it as the base
@@ -92,54 +91,26 @@ class FluxKontextModel(BaseModel):
                 base_model_path = model_path
 
         self.print_and_status_update("Loading transformer")
-        transformer = FluxTransformer2DModel.from_pretrained(
-            transformer_path,
-            subfolder=transformer_subfolder,
-            torch_dtype=dtype
+        transformer = FluxTransformer2DModel.load(
+            model_path, **self.component_load_kwargs("transformer")
         )
-        transformer.to(self.quantize_device, dtype=dtype)
-
-        if self.model_config.quantize:
-            # patch the state dict method
-            patch_dequantization_on_save(transformer)
-            quantization_type = get_qtype(self.model_config.qtype)
-            self.print_and_status_update("Quantizing transformer")
-            quantize(transformer, weights=quantization_type,
-                     **self.model_config.quantize_kwargs)
-            freeze(transformer)
-            transformer.to(self.device_torch)
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
 
         flush()
 
         self.print_and_status_update("Loading T5")
-        tokenizer_2 = T5TokenizerFast.from_pretrained(
-            base_model_path, subfolder="tokenizer_2", torch_dtype=dtype
+        tokenizer_2 = T5TextEncoder.load_tokenizer(base_model_path)
+        text_encoder_2 = T5TextEncoder.load(
+            base_model_path, **self.component_load_kwargs("te")
         )
-        text_encoder_2 = T5EncoderModel.from_pretrained(
-            base_model_path, subfolder="text_encoder_2", torch_dtype=dtype
-        )
-        text_encoder_2.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing T5")
-            quantize(text_encoder_2, weights=get_qtype(
-                self.model_config.qtype))
-            freeze(text_encoder_2)
-            flush()
 
         self.print_and_status_update("Loading CLIP")
-        text_encoder = CLIPTextModel.from_pretrained(
-            base_model_path, subfolder="text_encoder", torch_dtype=dtype)
-        tokenizer = CLIPTokenizer.from_pretrained(
-            base_model_path, subfolder="tokenizer", torch_dtype=dtype)
-        text_encoder.to(self.device_torch, dtype=dtype)
+        text_encoder = CLIPTextEncoder.load_model(
+            base_model_path, dtype=dtype, device=self.device_torch
+        )
+        tokenizer = CLIPTextEncoder.load_tokenizer(base_model_path, use_fast=False)
 
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKL.from_pretrained(
-            base_model_path, subfolder="vae", torch_dtype=dtype)
+        vae = KLVAE.load_model(base_model_path, dtype=dtype)
 
         self.noise_scheduler = FluxKontextModel.get_train_scheduler()
 
@@ -166,11 +137,13 @@ class FluxKontextModel(BaseModel):
         pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        # just to make sure everything is on the right device and dtype
-        text_encoder[0].to(self.device_torch)
+        # low_vram: text encoders stay on cpu; get_prompt_embeds moves them
+        # to the gpu on demand
+        if not self.low_vram:
+            text_encoder[0].to(self.device_torch)
+            text_encoder[1].to(self.device_torch)
         text_encoder[0].requires_grad_(False)
         text_encoder[0].eval()
-        text_encoder[1].to(self.device_torch)
         text_encoder[1].requires_grad_(False)
         text_encoder[1].eval()
         pipe.transformer = pipe.transformer.to(self.device_torch)

extensions_built_in/diffusion_models/hidream/hidream_e1_model.py

@@ -7,7 +7,7 @@ from toolkit.accelerator import unwrap_model
 import torch
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.config_modules import GenerateImageConfig
-from diffusers.models import HiDreamImageTransformer2DModel
+from toolkit.models.v2.diffusion_models.hidream import HiDreamImageTransformer2DModel
 
 import torch.nn.functional as F
 from PIL import Image

extensions_built_in/diffusion_models/hidream/hidream_model.py

@@ -9,7 +9,11 @@ from toolkit import train_tools
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from PIL import Image
 from toolkit.models.base_model import BaseModel
-from diffusers import AutoencoderKL, TorchAoConfig
+from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
+from toolkit.models.v2.text_encoders.llama import LlamaTextEncoder
+from toolkit.models.v2.text_encoders.clip import CLIPTextEncoderWithProjection
+from toolkit.models.v2.vae.autoencoder_kl import KLVAE
+from diffusers import TorchAoConfig
 from toolkit.basic import flush
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
@@ -18,8 +22,6 @@ from toolkit.dequantize import patch_dequantization_on_save
 from toolkit.accelerator import get_accelerator, unwrap_model
 from optimum.quanto import freeze, QTensor
 from toolkit.util.mask import generate_random_mask, random_dialate_mask
-from toolkit.util.quantize import quantize, get_qtype
-from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer, TorchAoConfig as TorchAoConfigTransformers
 from .src.pipelines.hidream_image.pipeline_hidream_image import HiDreamImagePipeline
 from .src.models.transformers.transformer_hidream_image import HiDreamImageTransformer2DModel
 from .src.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler
@@ -28,15 +30,6 @@ from einops import rearrange, repeat
 import random
 import torch.nn.functional as F
 from tqdm import tqdm
-from transformers import (
-    CLIPTextModelWithProjection,
-    CLIPTokenizer,
-    T5EncoderModel,
-    T5Tokenizer,
-    LlamaForCausalLM,
-    PreTrainedTokenizerFast
-)
-
 if TYPE_CHECKING:
     from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
 
@@ -103,117 +96,61 @@ class HidreamModel(BaseModel):
             use_fast=False
         )
         
-        text_encoder_4 = LlamaForCausalLM.from_pretrained(
+        # load + quantize + offload + placement, all driven by model_config
+        text_encoder_4 = LlamaTextEncoder.load(
             llama_model_path,
+            subfolder="",
             output_hidden_states=True,
             output_attentions=True,
-            torch_dtype=torch.bfloat16,
+            **self.component_load_kwargs("te"),
         )
-        text_encoder_4.to(self.device_torch, dtype=dtype)
-        
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing llama 8b model")
-            quantization_type = get_qtype(self.model_config.qtype_te)
-            quantize(text_encoder_4, weights=quantization_type)
-            freeze(text_encoder_4)
-        
-        if self.low_vram:
-            # unload it for now
-            text_encoder_4.to('cpu')
-            
+
         flush()
-        
+
         self.print_and_status_update("Loading transformer")
-            
-        transformer = self.hidream_transformer_class.from_pretrained(
-            model_path, 
-            subfolder="transformer", 
-            torch_dtype=torch.bfloat16
+
+        transformer = self.hidream_transformer_class.load(
+            model_path, **self.component_load_kwargs("transformer")
         )
-        
-        if not self.low_vram:
-            transformer.to(self.device_torch, dtype=dtype)
-        
-        if self.model_config.quantize:
-            self.print_and_status_update("Quantizing transformer")
-            quantization_type = get_qtype(self.model_config.qtype)
-            if self.low_vram:
-                # move and quantize only certain pieces at a time.
-                all_blocks = list(transformer.double_stream_blocks) + list(transformer.single_stream_blocks)
-                self.print_and_status_update(" - quantizing transformer blocks")
-                for block in tqdm(all_blocks):
-                    block.to(self.device_torch, dtype=dtype)
-                    quantize(block, weights=quantization_type)
-                    freeze(block)
-                    block.to('cpu')
-                    # flush()
-                
-                self.print_and_status_update(" - quantizing extras")
-                transformer.to(self.device_torch, dtype=dtype)
-                quantize(transformer, weights=quantization_type)
-                freeze(transformer)
-            else: 
-                quantize(transformer, weights=quantization_type)
-                freeze(transformer)
-            
-        if self.low_vram:
-            # unload it for now
-            transformer.to('cpu')
-        
+
         flush()
         
         self.print_and_status_update("Loading vae")
         
-        vae = AutoencoderKL.from_pretrained(
-            extras_path,
-            subfolder="vae",
-            torch_dtype=torch.bfloat16
-        ).to(self.device_torch, dtype=dtype)
+        vae = KLVAE.load_model(extras_path, dtype=torch.bfloat16).to(
+            self.device_torch, dtype=dtype
+        )
         
         
         self.print_and_status_update("Loading clip encoders")
         
-        text_encoder = CLIPTextModelWithProjection.from_pretrained(
-            extras_path,
-            subfolder="text_encoder",
-            torch_dtype=torch.bfloat16
+        text_encoder = CLIPTextEncoderWithProjection.load_model(
+            extras_path, dtype=torch.bfloat16
         ).to(self.device_torch, dtype=dtype)
-        
-        tokenizer = CLIPTokenizer.from_pretrained(
-            extras_path,
-            subfolder="tokenizer"
+
+        tokenizer = CLIPTextEncoderWithProjection.load_tokenizer(
+            extras_path, use_fast=False
         )
-        
-        text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
-            extras_path,
-            subfolder="text_encoder_2",
-            torch_dtype=torch.bfloat16
+
+        text_encoder_2 = CLIPTextEncoderWithProjection.load_model(
+            extras_path, dtype=torch.bfloat16, subfolder="text_encoder_2"
         ).to(self.device_torch, dtype=dtype)
-        
-        tokenizer_2 = CLIPTokenizer.from_pretrained(
-            extras_path,
-            subfolder="tokenizer_2"
+
+        tokenizer_2 = CLIPTextEncoderWithProjection.load_tokenizer(
+            extras_path, subfolder="tokenizer_2", use_fast=False
         )
         
         flush()
         self.print_and_status_update("Loading T5 encoders")
         
-        text_encoder_3 = T5EncoderModel.from_pretrained(
-            extras_path,
-            subfolder="text_encoder_3",
-            torch_dtype=torch.bfloat16
-        ).to(self.device_torch, dtype=dtype)
+        # load + quantize + offload + placement, all driven by model_config
+        text_encoder_3 = T5TextEncoder.load(
+            extras_path, subfolder="text_encoder_3", **self.component_load_kwargs("te")
+        )
+        flush()
         
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing T5")
-            quantization_type = get_qtype(self.model_config.qtype_te)
-            quantize(text_encoder_3, weights=quantization_type)
-            freeze(text_encoder_3)
-            flush()
-        
-        tokenizer_3 = T5Tokenizer.from_pretrained(
-            extras_path,
-            subfolder="tokenizer_3"
+        tokenizer_3 = T5TextEncoder.load_tokenizer(
+            extras_path, subfolder="tokenizer_3", use_fast=False
         )
         flush()
         
@@ -432,21 +369,8 @@ class HidreamModel(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ['double_stream_blocks', 'single_stream_blocks']
 
-    def convert_lora_weights_before_save(self, state_dict):
-        # currently starte with transformer. but needs to start with diffusion_model. for comfyui
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        # saved as diffusion_model. but needs to be transformer. for ai-toolkit
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
     
     def get_base_model_version(self):
         return "hidream_i1"

extensions_built_in/diffusion_models/hidream/hidream_o1_model.py

@@ -1,4 +1,5 @@
 import os
+from toolkit.models.v2._mixin import OstrisTransformersMixin
 from typing import List, Optional
 
 import torch
@@ -15,8 +16,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 from safetensors.torch import load_file, save_file
 from toolkit.accelerator import unwrap_model
 from optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
-from toolkit.memory_management import MemoryManager
 
 from transformers import AutoProcessor
 from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
@@ -29,6 +28,17 @@ from .src.hidream_o1.model_config import model_config
 if TYPE_CHECKING:
     from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
 
+
+class HidreamO1Transformer(Qwen3VLForConditionalGeneration, OstrisTransformersMixin):
+    """The o1 DiT-in-LLM: the vendored Qwen3VL with the image-diffusion heads
+    (x_embedder / t_embedder1 / final_layer2). The generic Qwen3VLTextEncoder
+    must NOT be used here — it drops those keys as unexpected."""
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["model.language_model.layers"]
+
+
 scheduler_config = {
     "num_train_timesteps": 1000,
     "shift": 3.0,
@@ -130,7 +140,10 @@ class HidreamO1Model(BaseModel):
         self.use_old_lokr_format = False
         self.is_flow_matching = True
         self.is_transformer = True
-        self.target_lora_modules = ["Qwen3VLForConditionalGeneration"]
+        self.target_lora_modules = [
+            "Qwen3VLForConditionalGeneration",
+            "HidreamO1Transformer",
+        ]
         self.noise_scale = self.model_config.model_kwargs.get(
             "noise_scale", DEFAULT_NOISE_SCALE
         )
@@ -189,7 +202,7 @@ class HidreamO1Model(BaseModel):
             )
 
             # transformer.load_state_dict(state_dict, assign=True)
-            transformer = Qwen3VLForConditionalGeneration.from_pretrained(
+            transformer = HidreamO1Transformer.from_pretrained(
                 None,
                 config=Qwen3VLConfig(**model_config),
                 state_dict=state_dict,
@@ -197,30 +210,13 @@ class HidreamO1Model(BaseModel):
             )
             del state_dict  # free memory
         else:
-            transformer = Qwen3VLForConditionalGeneration.from_pretrained(
+            transformer = HidreamO1Transformer.from_pretrained(
                 model_path,
                 torch_dtype=self.torch_dtype,
             )
         flush()
-        if not self.model_config.low_vram:
-            transformer.to(self.device_torch)
-
-        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,
-                ignore_modules=[],
-            )
-
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
 
         # move over to device now if low vram

extensions_built_in/diffusion_models/hidream/src/models/transformers/transformer_hidream_image.py

@@ -8,6 +8,8 @@ from einops import repeat
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
 from diffusers.models.modeling_utils import ModelMixin
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
 from diffusers.utils.torch_utils import maybe_allow_in_graph
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
@@ -228,9 +230,15 @@ class HiDreamImageBlock(nn.Module):
         )
 
 class HiDreamImageTransformer2DModel(
-    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin
 ):
     _supports_gradient_checkpointing = True
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["double_stream_blocks", "single_stream_blocks"]
+
     _no_split_modules = ["HiDreamImageBlock"]
 
     @register_to_config

extensions_built_in/diffusion_models/ideogram4/ideogram4.py

@@ -17,16 +17,19 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 )
 from toolkit.accelerator import unwrap_model
 from toolkit.metadata import get_meta_for_safetensors
-from toolkit.memory_management import MemoryManager
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
-from optimum.quanto import freeze, QTensor
+from optimum.quanto import QTensor
 
 import huggingface_hub
 from huggingface_hub.errors import EntryNotFoundError
 from transformers import AutoModel, AutoTokenizer
 
 from .src.transformer import Ideogram4Config, Ideogram4Transformer2DModel
-from .src.vae import AutoEncoder, AutoEncoderParams, convert_diffusers_state_dict
+from toolkit.models.v2.vae.flux2_kl import (
+    AutoEncoder,
+    AutoEncoderParams,
+    convert_diffusers_state_dict,
+)
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLModelEncoder
 from .src.latent_norm import get_latent_norm
 from .src.pipeline import (
     Ideogram4Pipeline,
@@ -220,8 +223,8 @@ class Ideogram4Model(BaseModel):
         self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}")
 
         tokenizer = AutoTokenizer.from_pretrained(te_path, token=HF_TOKEN)
-        text_encoder = AutoModel.from_pretrained(
-            te_path, torch_dtype=dtype, token=HF_TOKEN
+        text_encoder = Qwen3VLModelEncoder.load_model(
+            te_path, dtype=dtype, subfolder="", token=HF_TOKEN
         )
         flush()
 
@@ -234,9 +237,6 @@ class Ideogram4Model(BaseModel):
         self.print_and_status_update("Loading transformer")
 
         transformer_config = Ideogram4Config()
-        with torch.device("meta"):
-            transformer = Ideogram4Transformer2DModel(transformer_config)
-
         self.print_and_status_update("  - fetching transformer weights")
         state_dict = _load_component_state_dict(
             base, "transformer", "diffusion_pytorch_model"
@@ -246,7 +246,9 @@ class Ideogram4Model(BaseModel):
             state_dict, dtype, self.device_torch, self.model_config.low_vram
         )
         self.print_and_status_update("  - loading transformer state dict")
-        transformer.load_state_dict(state_dict, assign=True)
+        transformer = Ideogram4Transformer2DModel.load_from_state_dict(
+            state_dict, dtype, config=transformer_config
+        )
         del state_dict
         flush()
 
@@ -265,8 +267,7 @@ class Ideogram4Model(BaseModel):
         self.print_and_status_update("Loading VAE")
         vae_sd = _load_component_state_dict(base, "vae", "diffusion_pytorch_model")
         vae_sd = convert_diffusers_state_dict(vae_sd)
-        vae = AutoEncoder(AutoEncoderParams())
-        vae.load_state_dict(vae_sd)
+        vae = AutoEncoder.load_from_state_dict(vae_sd, self.vae_torch_dtype)
         del vae_sd
         vae.to(self.vae_device_torch, dtype=dtype)
         vae.eval()
@@ -362,58 +363,13 @@ class Ideogram4Model(BaseModel):
 
         transformer = self._load_transformer(base)
 
-        if self.model_config.quantize:
-            self.print_and_status_update("Quantizing Transformer")
-            quantize_model(self, transformer)
-            flush()
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
-        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,
-                ignore_modules=[
-                    transformer.rotary_emb.inv_freq,
-                    transformer.input_proj,
-                    transformer.llm_cond_proj,
-                ],
-            )
-        elif self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
-        else:
-            # quantize_model leaves the model on CPU; make sure it lands on device.
-            transformer.to(self.device_torch)
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
 
         tokenizer, text_encoder = self._load_text_encoder(base)
-        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,
-            )
-        elif self.model_config.low_vram:
-            self.print_and_status_update("Moving text encoder to CPU")
-            text_encoder.to("cpu")
-        else:
-            self.print_and_status_update("Moving text encoder to device")
-            text_encoder.to(self.device_torch)
+        # quantize + offload + placement, all driven by model_config
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
         flush()
 
         vae = self._load_vae(base)
@@ -620,16 +576,5 @@ class Ideogram4Model(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["layers"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd

extensions_built_in/diffusion_models/ideogram4/src/transformer.py

@@ -12,6 +12,8 @@ import math
 from dataclasses import dataclass
 
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint
@@ -355,9 +357,20 @@ class Ideogram4FinalLayer(nn.Module):
         return self.linear(self.norm_final(x) * scale)
 
 
-class Ideogram4Transformer2DModel(nn.Module):
+class Ideogram4Transformer2DModel(nn.Module, OstrisModelMixin):
     """Ideogram 4 flow-matching transformer."""
 
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]
+
+    def get_offload_ignore_modules(self):
+        return [
+            self.rotary_emb.inv_freq,
+            self.input_proj,
+            self.llm_cond_proj,
+        ]
+
     def __init__(self, config: Ideogram4Config) -> None:
         super().__init__()
         self.config = config

extensions_built_in/diffusion_models/krea2/krea2.py

@@ -25,18 +25,21 @@ from safetensors.torch import load_file, save_file
 
 import huggingface_hub
 from huggingface_hub.errors import EntryNotFoundError
-from diffusers import AutoencoderKLQwenImage
 from transformers import (
     AutoProcessor,
     AutoTokenizer,
     Qwen2TokenizerFast,
-    Qwen3VLForConditionalGeneration,
 )
 from optimum.quanto import freeze
 
 from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig
 from toolkit.lora_special import LoRASpecialNetwork
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin
+from toolkit.models.v2.text_encoders.qwen3_vl import (
+    Qwen3VLTextEncoder,
+    patch_qwen_vl_patch_embed,
+)
 from toolkit.basic import flush
 from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import (
@@ -44,7 +47,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 )
 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.mmdit import (
@@ -100,30 +102,6 @@ QWEN_IMAGE_VAE_PATH = "Qwen/Qwen-Image"
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 
 
-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
-    Qwen3VLCaptioner extension.)"""
-    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 _load_mmdit_state_dict(name_or_path: str, filename: Optional[str]) -> dict:
     """Load the MMDiT weights from a local safetensors file/dir or the HF hub.
 
@@ -163,7 +141,7 @@ def _load_mmdit_state_dict(name_or_path: str, filename: Optional[str]) -> dict:
     return load_file(path)
 
 
-class Krea2Model(BaseModel):
+class Krea2Model(QwenImageVAEHolderMixin, BaseModel):
     arch = "krea2"
 
     def __init__(
@@ -238,21 +216,15 @@ class Krea2Model(BaseModel):
         mmdit_kwargs.update(self.model_config.model_kwargs.get("mmdit_config", {}))
         config = SingleMMDiTConfig(**mmdit_kwargs)
 
-        # Build on meta, then materialize straight from the checkpoint.
-        with torch.device("meta"):
-            transformer = SingleStreamDiT(config)
-
         self.print_and_status_update("  - fetching transformer weights")
         state_dict = _load_mmdit_state_dict(
             self.model_config.name_or_path,
             self.model_config.model_kwargs.get("checkpoint_filename", None),
         )
-        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)
+        transformer = SingleStreamDiT.load_from_state_dict(
+            state_dict, dtype, config=config
+        )
         del state_dict
         flush()
         return transformer
@@ -268,8 +240,8 @@ class Krea2Model(BaseModel):
         processor = Qwen2TokenizerFast.from_pretrained(
             te_path, max_length=self.max_text_length, token=HF_TOKEN
         )
-        text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
-            te_path, torch_dtype=dtype, token=HF_TOKEN
+        text_encoder = Qwen3VLTextEncoder.load_model(
+            te_path, dtype=dtype, subfolder="", token=HF_TOKEN
         )
         vl_processor = None
         if self.is_edit:
@@ -281,8 +253,7 @@ class Krea2Model(BaseModel):
         else:
             # We only ever encode text, so the vision tower is dead weight -- drop it to
             # free VRAM and skip loading its (bf16-slow) Conv3d patch_embed onto the GPU.
-            if getattr(text_encoder.model, "visual", None) is not None:
-                text_encoder.model.visual = None
+            text_encoder.drop_vision_tower()
         text_encoder.eval()
         text_encoder.requires_grad_(False)
         flush()
@@ -291,8 +262,8 @@ class Krea2Model(BaseModel):
     def _load_vae(self):
         vae_path = self.model_config.model_kwargs.get("vae_path", QWEN_IMAGE_VAE_PATH)
         self.print_and_status_update(f"Loading Qwen-Image VAE from {vae_path}")
-        vae = AutoencoderKLQwenImage.from_pretrained(
-            vae_path, subfolder="vae", torch_dtype=self.vae_torch_dtype, token=HF_TOKEN
+        vae = QwenImageVAE.load_model(
+            vae_path, dtype=self.vae_torch_dtype, token=HF_TOKEN
         )
         vae.eval()
         vae.requires_grad_(False)
@@ -411,55 +382,12 @@ class Krea2Model(BaseModel):
             if self.model_config.qtype == "qfloat8":
                 self.model_config.qtype = "float8"
 
-        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,
-                ignore_modules=[
-                    module
-                    for module in transformer.modules()
-                    if isinstance(module, (SimpleModulation, DoubleSharedModulation))
-                ],
-            )
-
-        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)
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
 
         tokenizer, processor, 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)
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
         flush()
 
         vae = self._load_vae()
@@ -772,76 +700,9 @@ class Krea2Model(BaseModel):
         return False
 
     # ------------------------------------------------------------------
-    # VAE (Qwen-Image AutoencoderKLQwenImage -- same handling as qwen_image arch)
+    # VAE (Qwen-Image AutoencoderKLQwenImage -- shared QwenImageVAEHolderMixin)
     # ------------------------------------------------------------------
-    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()
-        self.vae.requires_grad_(False)
-
-        image_list = [image.to(device, dtype=dtype) for image in image_list]
-        images = torch.stack(image_list).to(device, dtype=dtype)
-
-        # AutoencoderKLQwenImage is a video VAE: add a frame dim.
-        images = images.unsqueeze(2)
-        latents = self.vae.encode(images).latent_dist.sample()
-
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.vae.config.z_dim, 1, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
-            1, self.vae.config.z_dim, 1, 1, 1
-        ).to(latents.device, latents.dtype)
-
-        latents = (latents - latents_mean) * latents_std
-        latents = latents.squeeze(2)  # drop frame dim
-        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(device, dtype=dtype)
-        latents = latents.unsqueeze(2)  # add frame dim
-
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.vae.config.z_dim, 1, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents_std = (
-            torch.tensor(self.vae.config.latents_std)
-            .view(1, self.vae.config.z_dim, 1, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents = latents * latents_std + latents_mean
-
-        # Full-resolution decode spikes VRAM; tile it when low on VRAM (decode
-        # only -- encode stays untiled).
-        tiled = self.model_config.low_vram
-        if tiled:
-            self.vae.enable_tiling()
-        try:
-            images = self.vae.decode(latents).sample
-        finally:
-            if tiled:
-                self.vae.disable_tiling()
-        images = images.squeeze(2)  # drop frame dim
-        return images.to(device, dtype=dtype)
-
+    vae_decode_tiled_on_low_vram = True
     # ------------------------------------------------------------------
     # Saving / bookkeeping
     # ------------------------------------------------------------------
@@ -867,14 +728,5 @@ class Krea2Model(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["blocks"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        return {
-            k.replace("transformer.", "diffusion_model."): v
-            for k, v in state_dict.items()
-        }
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        return {
-            k.replace("diffusion_model.", "transformer."): v
-            for k, v in state_dict.items()
-        }

extensions_built_in/diffusion_models/krea2/src/mmdit.py

@@ -20,6 +20,8 @@ import math
 from dataclasses import dataclass
 
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
@@ -408,7 +410,15 @@ class SingleStreamBlock(nn.Module):
         return x
 
 
-class SingleStreamDiT(nn.Module):
+class SingleStreamDiT(nn.Module, OstrisModelMixin):
+    def get_offload_ignore_modules(self):
+        # modulation modules hold tiny live state the offloader must not page
+        return [
+            module
+            for module in self.modules()
+            if isinstance(module, (SimpleModulation, DoubleSharedModulation))
+        ]
+
     def __init__(self, config: SingleMMDiTConfig):
         super().__init__()
         self.config = config

extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py

@@ -7,12 +7,16 @@ from typing import Any, Dict, Tuple
 import torch
 from accelerate import init_empty_weights
 
-from diffusers import (
-    AutoencoderKLLTX2Audio,
-    AutoencoderKLLTX2Video,
+from toolkit.models.v2.diffusion_models.ltx2 import (
+    LTX2TextConnectors,
     LTX2VideoTransformer3DModel,
+    LTX2Vocoder,
+    LTX2VocoderWithBWE,
+)
+from toolkit.models.v2.vae.ltx2 import (
+    LTX2AudioVAE as AutoencoderKLLTX2Audio,
+    LTX2VideoVAE as AutoencoderKLLTX2Video,
 )
-from diffusers.pipelines.ltx2 import LTX2TextConnectors, LTX2Vocoder, LTX2VocoderWithBWE
 from diffusers.utils.import_utils import is_accelerate_available
 
 

extensions_built_in/diffusion_models/ltx2/ltx2.py

@@ -17,28 +17,30 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 )
 from accelerate import init_empty_weights
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
-from toolkit.memory_management import MemoryManager
 from toolkit.paths import MODELS_PATH
+from toolkit.util.mixed_precision import attach_per_op_casting, pin_stored_fp32
 from safetensors.torch import load_file
 from PIL import Image
 import huggingface_hub
 
 try:
     from diffusers import LTX2Pipeline, LTX2ImageToVideoPipeline
-    from diffusers.models.autoencoders import (
-        AutoencoderKLLTX2Audio,
-        AutoencoderKLLTX2Video,
-    )
-    from diffusers.models.transformers import LTX2VideoTransformer3DModel
     from diffusers.pipelines.ltx2.export_utils import encode_video
     from transformers import (
         Gemma3ForConditionalGeneration,
         GemmaTokenizerFast,
     )
-    from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder, LTX2VocoderWithBWE
-    from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
+    from toolkit.models.v2.diffusion_models.ltx2 import (
+        LTX2TextConnectors,
+        LTX2VideoTransformer3DModel,
+        LTX2Vocoder,
+        LTX2VocoderWithBWE,
+    )
+    from toolkit.models.v2.vae.ltx2 import (
+        LTX2AudioVAE as AutoencoderKLLTX2Audio,
+        LTX2VideoVAE as AutoencoderKLLTX2Video,
+    )
+    from toolkit.models.v2.text_encoders.gemma3 import Gemma3TextEncoder
     from .convert_ltx2_to_diffusers import (
         get_model_state_dict_from_combined_ckpt,
         convert_ltx2_transformer,
@@ -293,48 +295,17 @@ class LTX2Model(BaseModel):
             )
             transformer = transformer.to(dtype)
         else:
-            transformer_path = model_path
-            transformer_subfolder = "transformer"
-            if os.path.exists(transformer_path):
-                transformer_subfolder = None
-                transformer_path = os.path.join(transformer_path, "transformer")
+            if os.path.exists(model_path):
                 # check if the path is a full checkpoint.
                 te_folder_path = os.path.join(model_path, "text_encoder")
                 # if we have the te, this folder is a full checkpoint, use it as the base
                 if os.path.exists(te_folder_path):
                     base_model_path = model_path
 
-            transformer = LTX2VideoTransformer3DModel.from_pretrained(
-                transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype
-            )
+            transformer = LTX2VideoTransformer3DModel.load_model(model_path, dtype=dtype)
 
-        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
-        ):
-            ignore_modules = []
-            for block in transformer.transformer_blocks:
-                ignore_modules.append(block.scale_shift_table)
-                ignore_modules.append(block.audio_scale_shift_table)
-                ignore_modules.append(block.video_a2v_cross_attn_scale_shift_table)
-                ignore_modules.append(block.audio_a2v_cross_attn_scale_shift_table)
-            ignore_modules.append(transformer.scale_shift_table)
-            ignore_modules.append(transformer.audio_scale_shift_table)
-            MemoryManager.attach(
-                transformer,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_transformer_percent,
-                ignore_modules=ignore_modules,
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
 
         flush()
 
@@ -347,7 +318,7 @@ class LTX2Model(BaseModel):
             tokenizer = GemmaTokenizerFast.from_pretrained(base_te_path)
 
             with init_empty_weights():
-                text_encoder = Gemma3ForConditionalGeneration(
+                text_encoder = Gemma3TextEncoder(
                     Gemma3Config(
                         **{
                             "boi_token_index": 255999,
@@ -420,47 +391,30 @@ class LTX2Model(BaseModel):
             tokenizer = GemmaTokenizerFast.from_pretrained(
                 self.model_config.te_name_or_path
             )
-            text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
-                self.model_config.te_name_or_path, dtype=dtype
+            text_encoder = Gemma3TextEncoder.load_model(
+                self.model_config.te_name_or_path, dtype=dtype, subfolder=""
             )
         elif self.ltx_te_path is not None:
             # pull from model specific te
             tokenizer = GemmaTokenizerFast.from_pretrained(self.ltx_te_path)
-            text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
-                self.ltx_te_path, dtype=dtype
+            text_encoder = Gemma3TextEncoder.load_model(
+                self.ltx_te_path, dtype=dtype, subfolder=""
             )
         else:
             # using combo hf repo
             tokenizer = GemmaTokenizerFast.from_pretrained(
                 self.model_config.name_or_path, subfolder="tokenizer"
             )
-            text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
-                self.model_config.name_or_path, subfolder="text_encoder", dtype=dtype
+            text_encoder = Gemma3TextEncoder.load_model(
+                self.model_config.name_or_path, dtype=dtype
             )
 
         # remove the vision tower
         text_encoder.model.vision_tower = None
         flush()
         
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            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,
-                ignore_modules=[
-                    text_encoder.model.language_model.base_model.embed_tokens
-                ],
-            )
-
+        # quantize + offload + placement, all driven by model_config
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
         text_encoder.to(self.device_torch, dtype=dtype)
         flush()
 
@@ -497,24 +451,16 @@ class LTX2Model(BaseModel):
             del combined_state_dict
             flush()
         else:
-            vae = AutoencoderKLLTX2Video.from_pretrained(
-                base_model_path, subfolder="vae", torch_dtype=dtype
-            )
-            audio_vae = AutoencoderKLLTX2Audio.from_pretrained(
-                base_model_path, subfolder="audio_vae", torch_dtype=dtype
-            )
+            vae = AutoencoderKLLTX2Video.load_model(base_model_path, dtype=dtype)
+            audio_vae = AutoencoderKLLTX2Audio.load_model(base_model_path, dtype=dtype)
 
-            connectors = LTX2TextConnectors.from_pretrained(
-                base_model_path, subfolder="connectors", torch_dtype=dtype
-            )
+            connectors = LTX2TextConnectors.load_model(base_model_path, dtype=dtype)
 
             vocoder_cls = LTX2Vocoder
             if self.ltx_version in ("2.3", "2.5"):
                 vocoder_cls = LTX2VocoderWithBWE
 
-            vocoder = vocoder_cls.from_pretrained(
-                base_model_path, subfolder="vocoder", torch_dtype=dtype
-            )
+            vocoder = vocoder_cls.load_model(base_model_path, dtype=dtype)
 
         self.noise_scheduler = LTX2Model.get_train_scheduler()
 
@@ -544,8 +490,10 @@ class LTX2Model(BaseModel):
             pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        # just to make sure everything is on the right device and dtype
-        text_encoder[0].to(self.device_torch)
+        # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
+        # to the gpu on demand
+        if not self.low_vram:
+            text_encoder[0].to(self.device_torch)
         text_encoder[0].requires_grad_(False)
         text_encoder[0].eval()
         flush()
@@ -1191,23 +1139,17 @@ class LTX2Model(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["transformer_blocks"]
 
+    lora_keys_use_comfy_prefix = True
+
     def convert_lora_weights_before_save(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        new_sd = convert_lora_diffusers_to_original(new_sd, version=self.ltx_version)
-        return new_sd
+        state_dict = super().convert_lora_weights_before_save(state_dict)
+        return convert_lora_diffusers_to_original(state_dict, version=self.ltx_version)
 
     def convert_lora_weights_before_load(self, state_dict):
         state_dict = convert_lora_original_to_diffusers(
             state_dict, version=self.ltx_version
         )
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
+        return super().convert_lora_weights_before_load(state_dict)
 
 
 class LTX23Model(LTX2Model):
@@ -1239,88 +1181,30 @@ class LTX25Model(LTX2Model):
     ltx_te_path = None
 
     # ------------------------------------------------------------------
-    # ComfyUI-style file resolution (mirrors MinimaxH3Model)
+    # ComfyUI-style file resolution (toolkit/models/v2/resolver.py)
     # ------------------------------------------------------------------
-    @staticmethod
-    def _find_file_recursive(root_dir: str, filename: str) -> Optional[str]:
-        if not os.path.isdir(root_dir):
-            return None
-        for dirpath, dirnames, filenames in os.walk(root_dir):
-            dirnames.sort()
-            if filename in filenames:
-                return os.path.join(dirpath, filename)
-        return None
-
     def _resolve_comfy_file(self, component: str) -> str:
-        """Find a weight file at its local location, or download it there
-        when (and only when) it is missing.
-
-        Search order: model_kwargs override, the repo-relative path under
-        MODELS_PATH (diffusion_models/, text_encoders/, vae/), the bare
-        filename at the root, any subfolder of the component's category
-        folder (recursive), then the hub — downloaded to the repo-relative
-        path under MODELS_PATH.
-        """
-        override = self.model_config.model_kwargs.get(f"{component}_path", None)
-        if override is not None:
-            if not os.path.exists(override):
-                raise FileNotFoundError(
-                    f"model_kwargs.{component}_path does not exist: {override}"
-                )
-            return override
-
-        rel_path = COMFY_LTX25_FILES[component]
-        filename = os.path.basename(rel_path)
-        category = os.path.dirname(rel_path)
-        for rel in (rel_path, filename):
-            candidate = os.path.join(MODELS_PATH, rel)
-            if os.path.exists(candidate):
-                return candidate
-        found = self._find_file_recursive(os.path.join(MODELS_PATH, category), filename)
-        if found is not None:
-            return found
-
-        repo_id = COMFY_LTX25_REPO
-        name_or_path = self.model_config.name_or_path
-        if (
-            name_or_path
-            and not os.path.exists(name_or_path)
-            and not name_or_path.endswith(".safetensors")
-            and "/" in name_or_path
-        ):
-            repo_id = name_or_path
-        self.print_and_status_update(
-            f"Downloading {rel_path} from {repo_id} into {MODELS_PATH}"
+        from toolkit.models.v2.resolver import (
+            repo_id_from_name_or_path,
+            resolve_comfy_file,
         )
-        return huggingface_hub.hf_hub_download(
-            repo_id=repo_id, filename=rel_path, token=HF_TOKEN, local_dir=MODELS_PATH
+
+        return resolve_comfy_file(
+            COMFY_LTX25_FILES[component],
+            repo_id=repo_id_from_name_or_path(
+                self.model_config.name_or_path, COMFY_LTX25_REPO
+            ),
+            override_path=self.model_config.model_kwargs.get(
+                f"{component}_path", None
+            ),
+            hf_token=HF_TOKEN,
+            status_fn=self.print_and_status_update,
         )
 
     def _resolve_named_file(self, path: str, component: str) -> str:
-        """Resolve an explicit .safetensors path: local file, models-folder
-        file, or an 'org/repo/path/file.safetensors' hub path (downloaded
-        into the models folder)."""
-        if os.path.exists(path):
-            return path
-        splits = path.split("/")
-        if len(splits) < 3:
-            raise ValueError(
-                f"Invalid {component} path: {path}. Must be a local file or "
-                "'repo_id/repo/filename.safetensors' to download from the Hugging Face Hub."
-            )
-        rel_path = "/".join(splits[2:])
-        for candidate in (
-            os.path.join(MODELS_PATH, rel_path),
-            os.path.join(MODELS_PATH, splits[-1]),
-        ):
-            if os.path.exists(candidate):
-                return candidate
-        return huggingface_hub.hf_hub_download(
-            repo_id="/".join(splits[:2]),
-            filename=rel_path,
-            token=HF_TOKEN,
-            local_dir=MODELS_PATH,
-        )
+        from toolkit.models.v2.resolver import resolve_named_file
+
+        return resolve_named_file(path, component=component, hf_token=HF_TOKEN)
 
     def _resolve_dit_path(self) -> str:
         name_or_path = self.model_config.name_or_path
@@ -1343,7 +1227,7 @@ class LTX25Model(LTX2Model):
         from the state dict. Works unchanged for bf16 checkpoints, where no
         quant markers exist and everything strict-loads."""
         from toolkit.util.comfy_quant_import import import_comfy_quantized_layers
-        from toolkit.util.ostris_quant import OstrisLinear
+        from toolkit.models.v2._mixin import OstrisModelMixin
 
         state_dict, num_quantized = import_comfy_quantized_layers(
             module, state_dict, orig_dtype=self.torch_dtype
@@ -1352,32 +1236,8 @@ class LTX25Model(LTX2Model):
             self.print_and_status_update(
                 f" - attached {num_quantized} pre-quantized ConvRot layers to {name}"
             )
-        result = module.load_state_dict(state_dict, assign=True, strict=False)
-        # quantized linears hold their weight as backend buffers and had their
-        # bias assigned by the importer, so both report as "missing" here
-        quantized_param_keys = set()
-        for mod_name, m in module.named_modules():
-            if isinstance(m, OstrisLinear):
-                quantized_param_keys.add(f"{mod_name}.weight")
-                if m.bias is not None:
-                    quantized_param_keys.add(f"{mod_name}.bias")
-        bad_missing = [k for k in result.missing_keys if k not in quantized_param_keys]
-        if bad_missing or result.unexpected_keys:
-            raise ValueError(
-                f"LTX-2.5 {name} load mismatch: missing {bad_missing[:8]}, "
-                f"unexpected {result.unexpected_keys[:8]}"
-            )
-        # nothing may be left on the meta device (e.g. a bias the importer
-        # should have filled)
-        leftover_meta = [
-            param_name
-            for param_name, p in module.named_parameters()
-            if p.is_meta
-        ]
-        if leftover_meta:
-            raise ValueError(
-                f"LTX-2.5 {name} load left meta parameters: {leftover_meta[:8]}"
-            )
+        # whitelist for quantized weights + leftover-meta check
+        OstrisModelMixin._load_state_dict_with_quantized(module, state_dict)
         return num_quantized
 
     def _load_gemma4_text_encoder(self, te_path: str, te_state_dict: dict):
@@ -1387,7 +1247,7 @@ class LTX25Model(LTX2Model):
         dropped, matching how the Gemma-3 vision tower was discarded."""
         from safetensors import safe_open
         from transformers import Gemma4TextConfig
-        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel
+        from toolkit.models.v2.text_encoders.gemma3 import Gemma4TextEncoder
 
         with safe_open(te_path, framework="pt") as f:
             metadata = f.metadata() or {}
@@ -1397,7 +1257,7 @@ class LTX25Model(LTX2Model):
         }
 
         with init_empty_weights():
-            text_encoder = Gemma4TextModel(Gemma4TextConfig(**text_config))
+            text_encoder = Gemma4TextEncoder(Gemma4TextConfig(**text_config))
 
         def strip_model_prefix(key: str) -> str:
             return key[len("model.") :] if key.startswith("model.") else key
@@ -1464,40 +1324,19 @@ class LTX25Model(LTX2Model):
         del transformer_sd
         if num_quantized_dit == 0:
             transformer = transformer.to(dtype)
+        else:
+            # mixed-precision comfy file (fp32 scale_shift tables next to bf16
+            # weights): comfy-style per-op input casting, with the stored-fp32
+            # pieces pinned so parent .to(dtype) casts cannot downcast them
+            attach_per_op_casting(transformer)
+            pin_stored_fp32(transformer)
         flush()
 
-        if self.model_config.quantize:
-            if num_quantized_dit:
-                self.print_and_status_update(
-                    "Transformer is pre-quantized (ConvRot); skipping quantize"
-                )
-            else:
-                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
-        ):
-            ignore_modules = []
-            for block in transformer.transformer_blocks:
-                ignore_modules.append(block.scale_shift_table)
-                ignore_modules.append(block.audio_scale_shift_table)
-                ignore_modules.append(block.video_a2v_cross_attn_scale_shift_table)
-                ignore_modules.append(block.audio_a2v_cross_attn_scale_shift_table)
-            ignore_modules.append(transformer.scale_shift_table)
-            ignore_modules.append(transformer.audio_scale_shift_table)
-            MemoryManager.attach(
-                transformer,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_transformer_percent,
-                ignore_modules=ignore_modules,
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
+        if num_quantized_dit:
+            # pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize
+            transformer.aitk_is_quantized = True
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
 
         self.print_and_status_update("Loading connectors")
@@ -1510,6 +1349,9 @@ class LTX25Model(LTX2Model):
         del connectors_sd, dit_sd
         if num_quantized_connectors == 0:
             connectors = connectors.to(dtype)
+        else:
+            attach_per_op_casting(connectors)
+            pin_stored_fp32(connectors)
         flush()
 
         # ---- text encoder (Gemma-4 12B, single comfy file) ----
@@ -1521,32 +1363,11 @@ class LTX25Model(LTX2Model):
         del te_state_dict
         flush()
 
-        if self.model_config.quantize_te:
-            if num_quantized_te:
-                self.print_and_status_update(
-                    "Text encoder is pre-quantized (ConvRot); skipping quantize"
-                )
-            else:
-                self.print_and_status_update("Quantizing Text Encoder")
-                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
-        ):
-            # layer_scalar is a bare tensor buffer on each decoder layer; the
-            # manager never enumerates it, so it must ride along explicitly
-            ignore_modules = [text_encoder.embed_tokens]
-            for layer in text_encoder.layers:
-                ignore_modules.append(layer.layer_scalar)
-            MemoryManager.attach(
-                text_encoder,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_text_encoder_percent,
-                ignore_modules=ignore_modules,
-            )
+        if num_quantized_te:
+            # pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize
+            text_encoder.aitk_is_quantized = True
+        # quantize + offload + placement, all driven by model_config
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
 
         text_encoder.to(self.device_torch)
         text_encoder.requires_grad_(False)

extensions_built_in/diffusion_models/mageflow/mageflow.py

@@ -34,8 +34,8 @@ 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 transformers import AutoProcessor, AutoTokenizer
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
 
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.models.base_model import BaseModel
@@ -46,8 +46,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 )
 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
@@ -177,22 +175,13 @@ class MageFlowModel(BaseModel):
         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 = MageFlow.load_from_state_dict(state_dict, dtype, config=params)
+        # rebuild the plain-tensor RoPE tables at full precision
         transformer.reset_rope()
         del state_dict
         flush()
@@ -209,8 +198,11 @@ class MageFlowModel(BaseModel):
         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
+        text_encoder = Qwen3VLTextEncoder.load_model(
+            te_path,
+            dtype=dtype,
+            subfolder=te_kwargs.get("subfolder", ""),
+            token=HF_TOKEN,
         )
         vl_processor = None
         if self.is_edit:
@@ -224,8 +216,7 @@ class MageFlowModel(BaseModel):
         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.drop_vision_tower()
         text_encoder.eval()
         text_encoder.requires_grad_(False)
         flush()
@@ -236,8 +227,8 @@ class MageFlowModel(BaseModel):
         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,
+        vae = MageVAE.load_model(
+            vae_path,
             sample_posterior=bool(
                 self.model_config.model_kwargs.get("vae_sample_posterior", True)
             ),
@@ -268,50 +259,13 @@ class MageFlowModel(BaseModel):
 
         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)
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         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)
+        # quantize + offload + placement, all driven by model_config
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
         flush()
 
         vae = self._load_vae()
@@ -629,18 +583,7 @@ class MageFlowModel(BaseModel):
     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()
-        }
-
+    lora_keys_use_comfy_prefix = True
 
 class MageFlowEditModel(MageFlowModel):
     arch = "mageflow_edit"

extensions_built_in/diffusion_models/mageflow/src/text_encoder.py

@@ -18,7 +18,8 @@ import math
 from typing import List, Optional
 
 import torch
-import torch.nn.functional as F
+
+from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed
 
 
 # Prompt templates from the reference (mage_flow/models/utils.py). ``start_idx``
@@ -57,30 +58,6 @@ def edit_prompt_body(instruction: str, num_refs: int) -> str:
     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"]:

extensions_built_in/diffusion_models/mageflow/src/transformer.py

@@ -28,6 +28,8 @@ from dataclasses import dataclass
 from typing import Any
 
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 import torch.nn as nn
 from torch import Tensor
 from torch.utils.checkpoint import checkpoint
@@ -681,7 +683,27 @@ class MageFlowParams:
     patch_size: int = 1
 
 
-class MageFlow(nn.Module):
+class MageFlow(nn.Module, OstrisModelMixin):
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]
+
+    @classmethod
+    def get_quantization_exclude_modules(cls):
+        # sensitive modules kept in full precision (fnmatch patterns):
+        #   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 __init__(self, params: MageFlowParams):
         super().__init__()
         self.params = params

extensions_built_in/diffusion_models/mageflow/src/vae.py

@@ -17,6 +17,8 @@ import os
 from functools import lru_cache
 
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 import torch.nn as nn
 import torch.nn.functional as F
 
@@ -591,7 +593,7 @@ def _load_state_dict(ckpt_path: str):
     return state
 
 
-class MageVAE(nn.Module):
+class MageVAE(nn.Module, OstrisModelMixin):
     """
     Encode: DConvEncoder (one-step diffusion) → latent [B, 128, H/16, W/16]
     Decode: DConvDenoiser + CoD Decoder       → image  [B, 3, H, W] in [-1, 1]
@@ -600,20 +602,32 @@ class MageVAE(nn.Module):
     latent_channels = 128
     downsample_factor = 16
 
-    def __init__(self, ckpt_path: str, sample_posterior: bool = True):
+    def __init__(self, ckpt_path: str = None, 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)
+        if ckpt_path is not None:
+            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()
+            # 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()
+
+    @classmethod
+    def load_from_state_dict(cls, state_dict, dtype=None, **kwargs):
+        model = cls(
+            ckpt_path=None,
+            sample_posterior=kwargs.pop("sample_posterior", True),
+        )
+        model._load_encoder(state_dict, "<state dict>")
+        model._load_decoder(state_dict, "<state dict>")
+        model._freeze_adaln_cache()
+        return model
 
     def _load_encoder(self, sd, ckpt_path):
         prefix = "student.dconv_encoder."

extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py

@@ -49,22 +49,25 @@ from toolkit.accelerator import unwrap_model
 from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
 from toolkit.basic import flush
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
-from toolkit.memory_management import MemoryManager
 from toolkit.metadata import get_meta_for_safetensors
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
+from toolkit.models.v2.resolver import (
+    find_file_recursive,
+    repo_id_from_name_or_path,
+    resolve_comfy_file,
+)
 from toolkit.paths import MODELS_PATH
 from toolkit.util.comfy_quant_import import import_comfy_quantized_layers
 from toolkit.util.ostris_quant import OstrisLinear
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
-from toolkit.util.quantize import get_qtype, quantize, quantize_model
-from optimum.quanto import freeze
 
 from .src import packing
 
 packing_video_exts = [".mp4", ".avi", ".mov", ".webm", ".mkv", ".wmv", ".m4v", ".flv"]
-from .src.audio_vae import MiniMaxH3AudioVAE, fold_audio_vae_weight_norm
+from .src.audio_vae import MiniMaxH3AudioVAE
 from .src.packing import (
     KEYFRAME_ENCODE_SEED,
     KEYFRAME_NOISE_AUG_T,
@@ -233,64 +236,23 @@ class MinimaxH3Model(BaseModel):
     # ------------------------------------------------------------------
     # Loading
     # ------------------------------------------------------------------
-    @staticmethod
-    def _find_file_recursive(root_dir: str, filename: str) -> Optional[str]:
-        """First (breadth-stable, sorted) match of ``filename`` anywhere under
-        ``root_dir``."""
-        if not os.path.isdir(root_dir):
-            return None
-        for dirpath, dirnames, filenames in os.walk(root_dir):
-            dirnames.sort()
-            if filename in filenames:
-                return os.path.join(dirpath, filename)
-        return None
-
     def _resolve_comfy_file(self, component: str) -> str:
-        """Find a weight file at its local location, or download it there
-        when (and only when) it is missing.
-
-        Search order: model_kwargs override, the repo-relative path under
-        MODELS_PATH (diffusion_models/, text_encoders/, vae/), the bare
-        filename at the root, any subfolder of the component's category
-        folder (recursive — e.g. diffusion_models/my_custom_sub/), the same
-        spots under name_or_path when it is a local folder, then the hub —
-        downloaded to the repo-relative path under MODELS_PATH.
-        """
-        override = self.model_config.model_kwargs.get(f"{component}_path", None)
-        if override is not None:
-            if not os.path.exists(override):
-                raise FileNotFoundError(
-                    f"model_kwargs.{component}_path does not exist: {override}"
-                )
-            return override
-
-        rel_path = COMFY_FILES[component]
-        filename = os.path.basename(rel_path)
-        category = os.path.dirname(rel_path)
-        roots = [MODELS_PATH]
+        """Find a weight file at its local location (model_kwargs override,
+        comfy layout under MODELS_PATH or a local name_or_path dir), or
+        download it there when (and only when) it is missing — see
+        toolkit/models/v2/resolver.py for the search order."""
         name_or_path = self.model_config.name_or_path
-        if name_or_path and os.path.isdir(name_or_path):
-            roots.append(name_or_path)
-        for root in roots:
-            for rel in (rel_path, filename):
-                candidate = os.path.join(root, rel)
-                if os.path.exists(candidate):
-                    return candidate
-        for root in roots:
-            found = self._find_file_recursive(os.path.join(root, category), filename)
-            if found is not None:
-                return found
-
-        import huggingface_hub
-
-        repo_id = COMFY_REPO
-        if name_or_path and not os.path.exists(name_or_path) and "/" in name_or_path:
-            repo_id = name_or_path
-        self.print_and_status_update(
-            f"Downloading {rel_path} from {repo_id} into {MODELS_PATH}"
+        extra_roots = (
+            [name_or_path] if name_or_path and os.path.isdir(name_or_path) else []
         )
-        return huggingface_hub.hf_hub_download(
-            repo_id=repo_id, filename=rel_path, local_dir=MODELS_PATH
+        return resolve_comfy_file(
+            COMFY_FILES[component],
+            repo_id=repo_id_from_name_or_path(name_or_path, COMFY_REPO),
+            override_path=self.model_config.model_kwargs.get(
+                f"{component}_path", None
+            ),
+            extra_roots=extra_roots,
+            status_fn=self.print_and_status_update,
         )
 
     def _dit_component(self) -> str:
@@ -322,7 +284,7 @@ class MinimaxH3Model(BaseModel):
         lora_path = self.model_config.assistant_lora_path
         if not os.path.exists(lora_path):
             filename = os.path.basename(lora_path)
-            found = self._find_file_recursive(
+            found = find_file_recursive(
                 os.path.join(MODELS_PATH, "loras"), filename
             )
             if found is not None:
@@ -409,46 +371,12 @@ class MinimaxH3Model(BaseModel):
         self.invert_assistant_lora = False
 
     def _load_transformer(self) -> MiniMaxH3Transformer:
-        dtype = self.torch_dtype
         dit_path = self._resolve_comfy_file(self._dit_component())
         self.print_and_status_update(f"Loading transformer from {dit_path}")
-        state_dict = load_file(dit_path)
-
-        params = MiniMaxH3TransformerParams()
-        table = state_dict.get("adaln_t_table", None)
-        if table is not None:
-            # pruned checkpoint: factored timestep table instead of the MLP
-            params.adaln_t_table_size = table.shape[0]
-            params.time_embed_dim = table.shape[1]
-
-        with torch.device("meta"):
-            transformer = MiniMaxH3Transformer(params)
-
-        # attach the pre-quantized (int8 ConvRot) linears onto the toolkit's
-        # quantization backends; the rest loads at its stored precision (the
-        # checkpoint's bf16/fp16/fp32 mix is deliberate)
-        state_dict, num_quantized = import_comfy_quantized_layers(
-            transformer, state_dict, orig_dtype=dtype
-        )
-        if num_quantized:
-            self.print_and_status_update(
-                f" - attached {num_quantized} pre-quantized ConvRot layers"
-            )
-        result = transformer.load_state_dict(state_dict, assign=True, strict=False)
-        quantized_weight_keys = {
-            f"{name}.weight"
-            for name, m in transformer.named_modules()
-            if isinstance(m, OstrisLinear)
-        }
-        bad_missing = [k for k in result.missing_keys if k not in quantized_weight_keys]
-        if bad_missing or result.unexpected_keys:
-            raise ValueError(
-                f"MiniMax-H3 transformer load mismatch: missing {bad_missing[:8]}, "
-                f"unexpected {result.unexpected_keys[:8]}"
-            )
-        del state_dict
-        flush()
-        return transformer
+        # the mixin single-file path: config sniffed from the checkpoint
+        # (adaln_t_table), pre-quantized ConvRot linears attached, everything
+        # else at its stored precision (the bf16/fp16/fp32 mix is deliberate)
+        return MiniMaxH3Transformer.load_model(dit_path, dtype=self.torch_dtype)
 
     def _load_text_encoder(self):
         from accelerate import init_empty_weights
@@ -474,7 +402,7 @@ class MinimaxH3Model(BaseModel):
             )
             config = AutoConfig.from_pretrained(te_path)
             config.text_config.num_hidden_layers = TEXT_ENCODER_LAYER
-            text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
+            text_encoder = Qwen3VLTextEncoder.from_pretrained(
                 te_path, config=config, torch_dtype=self.te_torch_dtype
             )
         else:
@@ -496,7 +424,7 @@ class MinimaxH3Model(BaseModel):
             config.text_config.num_hidden_layers = TEXT_ENCODER_LAYER
             config.tie_word_embeddings = False
             with init_empty_weights():
-                text_encoder = Qwen3VLForConditionalGeneration(config)
+                text_encoder = Qwen3VLTextEncoder(config)
             text_encoder.lm_head = None
 
             state_dict = load_file(te_file)
@@ -551,38 +479,9 @@ class MinimaxH3Model(BaseModel):
 
     def _load_vaes(self) -> MiniMaxH3VaeBundle:
         self.print_and_status_update("Loading video VAE")
-        video_sd = load_file(self._resolve_comfy_file("video_vae"))
-        # normalization stats ride along in the comfy file; the module holds
-        # them as non-persistent buffers, keep them float32
-        video_stats = {
-            k: video_sd.pop(k).float()
-            for k in ("latents_mean", "latents_std")
-            if k in video_sd
-        }
-        video_vae = MiniMaxH3VideoVAE()
-        video_vae.load_state_dict(video_sd, strict=True, assign=True)
-        for k, v in video_stats.items():
-            getattr(video_vae, k).copy_(v)
-        video_vae.eval().requires_grad_(False)
-        del video_sd
-
+        video_vae = MiniMaxH3VideoVAE.load_model(self._resolve_comfy_file("video_vae"))
         self.print_and_status_update("Loading audio VAE")
-        audio_sd = load_file(self._resolve_comfy_file("audio_vae"))
-        audio_stats = {
-            k: audio_sd.pop(k).float()
-            for k in ("latents_mean", "latents_std")
-            if k in audio_sd
-        }
-        # comfy repack ships the weight norm already folded; fold only if the
-        # raw parametrization is present (original-repo file)
-        if any(k.endswith("weight_g") for k in audio_sd.keys()):
-            audio_sd = fold_audio_vae_weight_norm(audio_sd)
-        audio_vae = MiniMaxH3AudioVAE()
-        audio_vae.load_state_dict(audio_sd, strict=True, assign=True)
-        for k, v in audio_stats.items():
-            getattr(audio_vae, k).copy_(v)
-        audio_vae.to(torch.float32).eval().requires_grad_(False)
-        del audio_sd
+        audio_vae = MiniMaxH3AudioVAE.load_model(self._resolve_comfy_file("audio_vae"))
         flush()
         return MiniMaxH3VaeBundle(video_vae, audio_vae)
 
@@ -596,54 +495,16 @@ class MinimaxH3Model(BaseModel):
         if self.model_config.assistant_lora_path is not None:
             self.load_training_adapter(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("Keeping transformer on CPU")
-            transformer.to("cpu")
-        else:
-            transformer.to(self.device_torch)
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
 
         tokenizer, processor, text_encoder = self._load_text_encoder()
-        te_prequantized = any(
-            isinstance(m, OstrisLinear) for m in text_encoder.modules()
-        )
-        if self.model_config.quantize_te and not te_prequantized:
-            self.print_and_status_update("Quantizing text encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
-        elif self.model_config.quantize_te:
-            self.print_and_status_update(
-                "Text encoder is already nvfp4/int8 quantized; skipping quantize_te"
-            )
-        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:
-            text_encoder.to("cpu")
-        else:
-            text_encoder.to(self.device_torch)
+        if any(isinstance(m, OstrisLinear) for m in text_encoder.modules()):
+            # already nvfp4/int8 quantized; aitk_post_load skips quantize_te
+            text_encoder.aitk_is_quantized = True
+        # quantize + offload + placement, all driven by model_config
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
         flush()
 
         vae_bundle = self._load_vaes()
@@ -1226,20 +1087,9 @@ class MinimaxH3Model(BaseModel):
             "*adaln_proj*",
         ]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        # ComfyUI's MiniMax-H3 keys are the original checkpoint keys, so the
-        # standard diffusion_model prefix maps directly
-        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()
-        }
-
+    # ComfyUI's MiniMax-H3 keys are the original checkpoint keys, so the
+    # standard diffusion_model prefix maps directly
+    lora_keys_use_comfy_prefix = True
 
 class MinimaxH3Ref2VAModel(MinimaxH3Model):
     """Reference-to-video (ref2va): the control images ride along as reference

extensions_built_in/diffusion_models/minimax_h3/src/audio_vae.py

@@ -17,6 +17,8 @@ weights::
 import math
 
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
@@ -420,7 +422,26 @@ class BigVGANDecoder(nn.Module):
         return torch.clamp(x, min=-1.0, max=1.0)
 
 
-class MiniMaxH3AudioVAE(nn.Module):
+class MiniMaxH3AudioVAE(nn.Module, OstrisModelMixin):
+    @classmethod
+    def load_from_state_dict(cls, state_dict, dtype=None, **kwargs):
+        """Comfy single-file load: stats buffers ride along; original-repo
+        files still carry the raw weight-norm parametrization and get folded."""
+        state_dict = dict(state_dict)
+        stats = {
+            k: state_dict.pop(k).float()
+            for k in ("latents_mean", "latents_std")
+            if k in state_dict
+        }
+        if any(k.endswith("weight_g") for k in state_dict.keys()):
+            state_dict = fold_audio_vae_weight_norm(state_dict)
+        model = cls()
+        model.load_state_dict(state_dict, strict=True, assign=True)
+        for k, v in stats.items():
+            getattr(model, k).copy_(v)
+        model.to(torch.float32).eval().requires_grad_(False)
+        return model
+
     """Frozen MiniMax-H3 audio autoencoder with diagonal latent normalization.
 
     - ``encode``: waveform ``(B, 1, samples)`` at 32 kHz -> latents ``(B, 32, T)``

extensions_built_in/diffusion_models/minimax_h3/src/transformer.py

@@ -30,6 +30,8 @@ import math
 from dataclasses import dataclass
 from typing import Optional, Tuple
 
+from toolkit.models.v2._mixin import OstrisModelMixin
+
 import torch
 import torch.nn.functional as F
 from torch import nn
@@ -356,7 +358,31 @@ class MiniMaxH3FinalLayer(nn.Module):
         return self.video_out(h), self.audio_out(h)
 
 
-class MiniMaxH3Transformer(nn.Module):
+class MiniMaxH3Transformer(nn.Module, OstrisModelMixin):
+    # comfy checkpoints carry a deliberate bf16/fp16/fp32 mix
+    aitk_cast_on_load = False
+
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict):
+        params = MiniMaxH3TransformerParams()
+        table = state_dict.get("adaln_t_table", None)
+        if table is not None:
+            # pruned checkpoint: factored timestep table instead of the MLP
+            params.adaln_t_table_size = table.shape[0]
+            params.time_embed_dim = table.shape[1]
+        return params
+
+    @classmethod
+    def aitk_from_config(cls, config):
+        from accelerate import init_empty_weights
+
+        with init_empty_weights(include_buffers=False):
+            return cls(config)
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["blocks"]
+
     def __init__(self, params: Optional[MiniMaxH3TransformerParams] = None):
         super().__init__()
         if params is None:

extensions_built_in/diffusion_models/minimax_h3/src/vae.py

@@ -22,6 +22,8 @@ import math
 from contextlib import nullcontext
 
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint
@@ -434,7 +436,24 @@ class ViTDecoder3d(nn.Module):
 # ---------------------------------------------------------------------------
 
 
-class MiniMaxH3VideoVAE(nn.Module):
+class MiniMaxH3VideoVAE(nn.Module, OstrisModelMixin):
+    @classmethod
+    def load_from_state_dict(cls, state_dict, dtype=None, **kwargs):
+        """Comfy single-file load: the normalization stats ride along in the
+        file and land in the non-persistent fp32 buffers."""
+        state_dict = dict(state_dict)
+        stats = {
+            k: state_dict.pop(k).float()
+            for k in ("latents_mean", "latents_std")
+            if k in state_dict
+        }
+        model = cls()
+        model.load_state_dict(state_dict, strict=True, assign=True)
+        for k, v in stats.items():
+            getattr(model, k).copy_(v)
+        model.eval().requires_grad_(False)
+        return model
+
     def __init__(
         self,
         in_channels=3,

extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py

@@ -6,21 +6,23 @@ import torch
 import yaml
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
 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 optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
-from toolkit.memory_management import MemoryManager
 
-from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+from transformers import Qwen3VLProcessor
 import torch.nn.functional as F
 
 try:
-    from diffusers import NucleusMoEImagePipeline, NucleusMoEImageTransformer2DModel, AutoencoderKLQwenImage
+    from diffusers import NucleusMoEImagePipeline, AutoencoderKLQwenImage
+    from toolkit.models.v2.diffusion_models.nucleus_image import (
+        NucleusMoEImageTransformer2DModel,
+    )
     from diffusers.models.transformers.transformer_nucleusmoe_image import SwiGLUExperts
 except ImportError:
     raise ImportError(
@@ -46,7 +48,7 @@ scheduler_config = {
 }
 
 
-class NucleusImageModel(BaseModel):
+class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel):
     arch = "nucleus_image"
 
     def __init__(
@@ -81,20 +83,14 @@ class NucleusImageModel(BaseModel):
 
         self.print_and_status_update("Loading transformer")
 
-        transformer_path = model_path
-        transformer_subfolder = "transformer"
-        if os.path.exists(transformer_path):
-            transformer_subfolder = None
-            transformer_path = os.path.join(transformer_path, "transformer")
+        if os.path.exists(model_path):
             # check if the path is a full checkpoint.
             te_folder_path = os.path.join(model_path, "text_encoder")
             # if we have the te, this folder is a full checkpoint, use it as the base
             if os.path.exists(te_folder_path):
                 base_model_path = model_path
 
-        transformer = NucleusMoEImageTransformer2DModel.from_pretrained(
-            transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype
-        )
+        transformer = NucleusMoEImageTransformer2DModel.load_model(model_path, dtype=dtype)
         
         # handle versions of pytorch that don't have grouped mm, by disabling it in the SwiGLUExperts
         if not hasattr(torch.nn.functional, "grouped_mm"):
@@ -102,26 +98,8 @@ class NucleusImageModel(BaseModel):
                 if isinstance(m, SwiGLUExperts):
                     m.use_grouped_mm = False
 
-        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,
-                ignore_modules=[
-                ],
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
 
         flush()
 
@@ -129,33 +107,14 @@ class NucleusImageModel(BaseModel):
         tokenizer = Qwen3VLProcessor.from_pretrained(
             base_model_path, subfolder="processor", torch_dtype=dtype
         )
-        text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
-            base_model_path, subfolder="text_encoder", torch_dtype=dtype
+        text_encoder = Qwen3VLTextEncoder.load(
+            base_model_path, **self.component_load_kwargs("te")
         )
 
-        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,
-            )
-
-        text_encoder.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
-
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKLQwenImage.from_pretrained(
-            base_model_path, subfolder="vae", torch_dtype=dtype
-        ).to(self.device_torch, dtype=dtype)
+        vae = QwenImageVAE.load_model(
+            base_model_path, dtype=dtype, device=self.device_torch
+        )
 
         self.noise_scheduler = NucleusImageModel.get_train_scheduler()
 
@@ -185,8 +144,10 @@ class NucleusImageModel(BaseModel):
             pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        # just to make sure everything is on the right device and dtype
-        text_encoder[0].to(self.device_torch)
+        # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
+        # to the gpu on demand
+        if not self.low_vram:
+            text_encoder[0].to(self.device_torch)
         text_encoder[0].requires_grad_(False)
         text_encoder[0].eval()
         flush()
@@ -214,41 +175,6 @@ class NucleusImageModel(BaseModel):
 
         return pipeline
 
-    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
-
-        # Move to vae to device if on cpu
-        if self.vae.device == torch.device("cpu"):
-            self.vae.to(device)
-        self.vae.eval()
-        self.vae.requires_grad_(False)
-        # move to device and dtype
-        image_list = [image.to(device, dtype=dtype) for image in image_list]
-        images = torch.stack(image_list).to(device, dtype=dtype)
-        # it uses wan vae, so add dim for frame count
-
-        images = images.unsqueeze(2)
-        latents = self.vae.encode(images).latent_dist.sample()
-
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.vae.config.z_dim, 1, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
-            1, self.vae.config.z_dim, 1, 1, 1
-        ).to(latents.device, latents.dtype)
-
-        latents = (latents - latents_mean) * latents_std
-        latents = latents.to(device, dtype=dtype)
-
-        latents = latents.squeeze(2)  # remove the frame count dimension
-
-        return latents
-
     def generate_single_image(
         self,
         pipeline: NucleusMoEImagePipeline,
@@ -405,16 +331,5 @@ class NucleusImageModel(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["transformer_blocks"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd

extensions_built_in/diffusion_models/omnigen2/__init__.py

@@ -1,4 +1,6 @@
 import os
+from toolkit.models.v2.text_encoders.qwen25_vl import Qwen25VLTextEncoder
+from toolkit.models.v2.vae.autoencoder_kl import KLVAE
 from typing import TYPE_CHECKING, List, Optional
 
 import torch
@@ -12,8 +14,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype
 from .src.pipelines.omnigen2.pipeline_omnigen2 import OmniGen2Pipeline
 from .src.models.transformers import OmniGen2Transformer2DModel
 from .src.models.transformers.repo import OmniGen2RotaryPosEmbed
@@ -78,48 +78,26 @@ class OmniGen2Model(BaseModel):
             extras_path, subfolder="processor", use_fast=True
         )
 
-        mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            extras_path, subfolder="mllm", torch_dtype=torch.bfloat16
+        # load + quantize + offload + placement, all driven by model_config
+        mllm = Qwen25VLTextEncoder.load(
+            extras_path, subfolder="mllm", **self.component_load_kwargs("te")
         )
-        mllm.to(self.device_torch, dtype=dtype)
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Qwen2.5 VL model")
-            quantization_type = get_qtype(self.model_config.qtype_te)
-            quantize(mllm, weights=quantization_type)
-            freeze(mllm)
-
-        if self.low_vram:
-            # unload it for now
-            mllm.to("cpu")
 
         flush()
 
         self.print_and_status_update("Loading transformer")
 
-        transformer = OmniGen2Transformer2DModel.from_pretrained(
-            model_path, subfolder="transformer", torch_dtype=torch.bfloat16
+        transformer = OmniGen2Transformer2DModel.load(
+            model_path, **self.component_load_kwargs("transformer")
         )
 
-        if not self.low_vram:
-            transformer.to(self.device_torch, dtype=dtype)
-
-        if self.model_config.quantize:
-            self.print_and_status_update("Quantizing transformer")
-            quantization_type = get_qtype(self.model_config.qtype)
-            quantize(transformer, weights=quantization_type)
-            freeze(transformer)
-
-        if self.low_vram:
-            # unload it for now
-            transformer.to("cpu")
-
         flush()
 
         self.print_and_status_update("Loading vae")
 
-        vae = AutoencoderKL.from_pretrained(
-            extras_path, subfolder="vae", torch_dtype=torch.bfloat16
-        ).to(self.device_torch, dtype=dtype)
+        vae = KLVAE.load_model(extras_path, dtype=torch.bfloat16).to(
+            self.device_torch, dtype=dtype
+        )
 
         flush()
         self.print_and_status_update("Loading Qwen2.5 VLProcessor")
@@ -343,21 +321,7 @@ class OmniGen2Model(BaseModel):
             return ["noise_refiner", "context_refiner", "ref_image_refiner", "layers"]
         return ["noise_refiner", "context_refiner", "layers"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        # currently starte with transformer. but needs to start with diffusion_model. for comfyui
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
-
-    def convert_lora_weights_before_load(self, state_dict):
-        # saved as diffusion_model. but needs to be transformer. for ai-toolkit
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
     def get_base_model_version(self):
         return "omnigen2"

extensions_built_in/diffusion_models/omnigen2/src/models/transformers/transformer_omnigen2.py

@@ -15,6 +15,8 @@ from diffusers.models.attention_processor import Attention
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.modeling_utils import ModelMixin
 
+from toolkit.models.v2._mixin import OstrisModelMixin
+
 from ..attention_processor import OmniGen2AttnProcessorFlash2Varlen, OmniGen2AttnProcessor
 from .repo import OmniGen2RotaryPosEmbed
 from .block_lumina2 import LuminaLayerNormContinuous, LuminaRMSNormZero, LuminaFeedForward, Lumina2CombinedTimestepCaptionEmbedding
@@ -177,7 +179,15 @@ class OmniGen2TransformerBlock(nn.Module):
         return hidden_states
 
 
-class OmniGen2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+class OmniGen2Transformer2DModel(
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["noise_refiner", "ref_image_refiner", "context_refiner", "layers"]
+
     """
     OmniGen2 Transformer 2D Model.
     

extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py

@@ -43,19 +43,20 @@ import torch
 import yaml
 
 from transformers import AutoTokenizer, Qwen3VLTextModel
-from optimum.quanto import freeze
+
 
 from toolkit.accelerator import unwrap_model
 from toolkit.basic import flush
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextOnlyEncoder
 from toolkit.models.FakeVAE import FakeVAE
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.train_tools import apply_noise_offset
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
+
 
 from .src.transformer_prx import PRXTransformer2DModel
 from .src.pipeline import PRXPixelPipeline
@@ -126,40 +127,22 @@ class PRXPixelT2IModel(BaseModel):
         self.print_and_status_update("Loading transformer")
         # from_pretrained reads config.json (bottleneck_size, resolution_embeds,
         # in_channels=3, ...) and the safetensors in one shot.
-        transformer = PRXTransformer2DModel.from_pretrained(
-            model_path, subfolder="transformer", torch_dtype=dtype
+        # load + quantize + offload + placement, all driven by model_config
+        transformer = PRXTransformer2DModel.load(
+            model_path, **self.component_load_kwargs("transformer")
         )
-        transformer.to(dtype=dtype)
-        flush()
-
-        if self.model_config.quantize:
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            flush()
-
-        if self.model_config.low_vram:
-            transformer.to("cpu")
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
         flush()
 
         # --- text encoder + tokenizer (Qwen3-VL text tower) ---
         self.print_and_status_update("Loading text encoder")
         tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer")
-        text_encoder = Qwen3VLTextModel.from_pretrained(
-            model_path, subfolder="text_encoder", torch_dtype=dtype
+        text_encoder = Qwen3VLTextOnlyEncoder.load(
+            model_path, **self.component_load_kwargs("te")
         )
-        text_encoder.to(self.te_device_torch)
         text_encoder.eval()
         text_encoder.requires_grad_(False)
         flush()
 
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing text encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
-
         # --- "VAE": identity, since PRXPixel is pixel space ---
         self.print_and_status_update("Preparing pixel-space VAE (identity)")
         vae = FakeVAE(scaling_factor=1.0)
@@ -332,14 +315,5 @@ class PRXPixelT2IModel(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["blocks"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        return {
-            k.replace("transformer.", "diffusion_model."): v
-            for k, v in state_dict.items()
-        }
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        return {
-            k.replace("diffusion_model.", "transformer."): v
-            for k, v in state_dict.items()
-        }

extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py

@@ -40,6 +40,8 @@ from diffusers.models.attention_dispatch import dispatch_attention_fn
 from diffusers.models.embeddings import get_timestep_embedding
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.modeling_utils import ModelMixin
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 from diffusers.models.normalization import RMSNorm
 
 
@@ -666,7 +668,7 @@ def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Te
     return seq
 
 
-class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
+class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin, OstrisModelMixin):
     r"""
     Transformer-based 2D model for text to image generation.
 
@@ -702,6 +704,12 @@ class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
             `PRXResolutionEmbedder`. Used by the PRX-7B variant.
     """
 
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["blocks"]
+
     config_name = "config.json"
     _supports_gradient_checkpointing = True
 

extensions_built_in/diffusion_models/qwen_image/qwen_image.py

@@ -7,28 +7,25 @@ from toolkit import train_tools
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from PIL import Image
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin
 from toolkit.basic import flush
 from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.accelerator import get_accelerator, unwrap_model
-from optimum.quanto import freeze, QTensor
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
 import torch.nn.functional as F
 from toolkit.memory_management import MemoryManager
+from toolkit.metadata import get_meta_for_safetensors
 from safetensors.torch import load_file
 
 from diffusers import (
     QwenImagePipeline,
-    QwenImageTransformer2DModel,
     AutoencoderKLQwenImage,
 )
-from transformers import (
-    Qwen2_5_VLForConditionalGeneration,
-    Qwen2Tokenizer,
-    Qwen2VLProcessor,
-)
+from transformers import Qwen2VLProcessor
+from toolkit.models.v2.diffusion_models.qwen_image import QwenImageTransformer2DModel
+from toolkit.models.v2.text_encoders.qwen25_vl import Qwen25VLTextEncoder
 from tqdm import tqdm
 from toolkit.util.qwen_vae_gradient_checkpointing import patch_qwen_vae_gradient_checkpointing
 
@@ -55,7 +52,7 @@ scheduler_config = {
 }
 
 
-class QwenImageModel(BaseModel):
+class QwenImageModel(QwenImageVAEHolderMixin, BaseModel):
     arch = "qwen_image"
     _qwen_image_keep_visual = False
     _qwen_pipeline = QwenImagePipeline
@@ -97,89 +94,33 @@ class QwenImageModel(BaseModel):
 
         self.print_and_status_update("Loading transformer")
 
-        if model_path.endswith(".safetensors"):
-            # load the safetensors file
-            transformer = QwenImageTransformer2DModel.from_single_file(
-                model_path,
-                config="Qwen/Qwen-Image",
-                subfolder="transformer",
-                torch_dtype=model_dtype,
-            )
-            transformer.to(model_dtype)
-
-        else:
-            transformer_path = model_path
-            transformer_subfolder = "transformer"
-            if os.path.exists(transformer_path):
-                transformer_subfolder = None
-                transformer_path = os.path.join(transformer_path, "transformer")
-                # check if the path is a full checkpoint.
-                te_folder_path = os.path.join(model_path, "text_encoder")
-                # if we have the te, this folder is a full checkpoint, use it as the base
-                if os.path.exists(te_folder_path):
-                    base_model_path = model_path
-
-            transformer = QwenImageTransformer2DModel.from_pretrained(
-                transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype
-            )
-
-        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")
+        if not model_path.endswith(".safetensors") and os.path.exists(model_path):
+            # check if the path is a full checkpoint.
+            te_folder_path = os.path.join(model_path, "text_encoder")
+            # if we have the te, this folder is a full checkpoint, use it as the base
+            if os.path.exists(te_folder_path):
+                base_model_path = model_path
 
+        # load + quantize + offload + placement, all driven by model_config
+        transformer = QwenImageTransformer2DModel.load(
+            model_path, **self.component_load_kwargs("transformer")
+        )
         flush()
 
         self.print_and_status_update("Text Encoder")
-        tokenizer = Qwen2Tokenizer.from_pretrained(
-            base_model_path, subfolder="tokenizer", torch_dtype=dtype
-        )
-        text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            base_model_path, subfolder="text_encoder", torch_dtype=dtype
-        )
+        tokenizer = Qwen25VLTextEncoder.load_tokenizer(base_model_path, use_fast=False)
+        text_encoder = Qwen25VLTextEncoder.load_model(base_model_path, dtype=dtype)
 
         # remove the visual model as it is not needed for image generation
+        # (before quantization, so the dead tower is never quantized)
         self.processor = None
         if not self._qwen_image_keep_visual:
-            text_encoder.model.visual = None
+            text_encoder.drop_vision_tower()
 
-        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,
-            )
-
-        text_encoder.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
 
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKLQwenImage.from_pretrained(
-            base_model_path, subfolder="vae", torch_dtype=dtype
-        )
+        vae = QwenImageVAE.load_model(base_model_path, dtype=dtype)
 
         self.noise_scheduler = QwenImageModel.get_train_scheduler()
 
@@ -220,8 +161,10 @@ class QwenImageModel(BaseModel):
             pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        # just to make sure everything is on the right device and dtype
-        text_encoder[0].to(self.device_torch)
+        # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
+        # to the gpu on demand
+        if not self.low_vram:
+            text_encoder[0].to(self.device_torch)
         text_encoder[0].requires_grad_(False)
         text_encoder[0].eval()
         flush()
@@ -393,17 +336,17 @@ class QwenImageModel(BaseModel):
         return False
 
     def save_model(self, output_path, meta, save_dtype):
-        # only save the unet
+        # comfy-format single-file save (diffusers keys ARE the comfy layout
+        # for qwen image); prequantized layers keep their quantized storage
         transformer: QwenImageTransformer2DModel = unwrap_model(self.model)
-        transformer.save_pretrained(
-            save_directory=os.path.join(output_path, "transformer"),
-            safe_serialization=True,
+        if not output_path.endswith(".safetensors"):
+            output_path += ".safetensors"
+        transformer.save_model(
+            output_path,
+            dtype=save_dtype,
+            metadata=get_meta_for_safetensors(meta, name=self.arch),
         )
 
-        meta_path = os.path.join(output_path, "aitk_meta.yaml")
-        with open(meta_path, "w") as f:
-            yaml.dump(meta, f)
-
     def get_loss_target(self, *args, **kwargs):
         noise = kwargs.get("noise")
         batch = kwargs.get("batch")
@@ -415,83 +358,5 @@ class QwenImageModel(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["transformer_blocks"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
-
-    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
-
-        # Move to vae to device if on cpu
-        if self.vae.device == torch.device("cpu"):
-            self.vae.to(device)
-        self.vae.eval()
-        self.vae.requires_grad_(False)
-        # move to device and dtype
-        image_list = [image.to(device, dtype=dtype) for image in image_list]
-        images = torch.stack(image_list).to(device, dtype=dtype)
-        # it uses wan vae, so add dim for frame count
-
-        images = images.unsqueeze(2)
-        latents = self.vae.encode(images).latent_dist.sample()
-
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.vae.config.z_dim, 1, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
-            1, self.vae.config.z_dim, 1, 1, 1
-        ).to(latents.device, latents.dtype)
-
-        latents = (latents - latents_mean) * latents_std
-        latents = latents.to(device, dtype=dtype)
-
-        latents = latents.squeeze(2)  # remove the frame count dimension
-
-        return latents
-
-    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(device, dtype=dtype)
-
-        # add frame count dim for wan vae
-        latents = latents.unsqueeze(2)
-
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.vae.config.z_dim, 1, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents_std = (
-            torch.tensor(self.vae.config.latents_std)
-            .view(1, self.vae.config.z_dim, 1, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents = latents * latents_std + latents_mean
-
-        images = self.vae.decode(latents).sample
-
-        images = images.squeeze(2)  # remove the frame count dimension
-
-        return images.to(device, dtype=dtype)

extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py

@@ -15,7 +15,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 )
 from toolkit.accelerator import get_accelerator, unwrap_model
 from optimum.quanto import freeze, QTensor
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
 import torch.nn.functional as F
 
 from diffusers import (

extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py

@@ -15,7 +15,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 )
 from toolkit.accelerator import get_accelerator, unwrap_model
 from optimum.quanto import freeze, QTensor
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
 import torch.nn.functional as F
 
 from diffusers import (

extensions_built_in/diffusion_models/wan22/wan22_14b_model.py

@@ -17,11 +17,12 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 )
 from toolkit.util.quantize import quantize_model
 from .wan22_pipeline import Wan22Pipeline
-from diffusers import WanTransformer3DModel
+from toolkit.models.v2.diffusion_models.wan import WanTransformer3DModel
 
 from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
 from torchvision.transforms import functional as TF
 
+from toolkit.metadata import get_meta_for_safetensors
 from toolkit.models.wan21.wan21 import Wan21
 from .wan22_5b_model import (
     scheduler_config,
@@ -94,6 +95,13 @@ class DualWanTransformer3DModel(torch.nn.Module):
     def dtype(self) -> torch.dtype:
         return self.torch_dtype
 
+    def get_offload_ignore_modules(self):
+        # both DiTs' fp32 modulation tables stay resident on the compute device
+        return (
+            self.transformer_1.get_offload_ignore_modules()
+            + self.transformer_2.get_offload_ignore_modules()
+        )
+
     @property
     def config(self):
         return self.transformer_1.config
@@ -127,17 +135,26 @@ class DualWanTransformer3DModel(torch.nn.Module):
             else:
                 t_name = "transformer_2"
 
+            # memory-managed transformers own their placement (layers bounce
+            # from cpu per forward); whole-model swaps would haul the full
+            # 14B up and defeat the offloading
+            managed = (
+                hasattr(self, "_memory_manager")
+                or hasattr(self.transformer_1, "_memory_manager")
+                or hasattr(self.transformer_2, "_memory_manager")
+            )
+
             # check if we are changing the active transformer, if so, we need to swap the one in
             # vram if low_vram is enabled
             # todo swap the loras as well
             if t_name != self._active_transformer_name:
-                if self.low_vram:
+                if self.low_vram and not managed:
                     getattr(self, self._active_transformer_name).to("cpu")
                     getattr(self, t_name).to(self.device_torch)
                     torch.cuda.empty_cache()
                 self._active_transformer_name = t_name
 
-        if self.transformer.device != hidden_states.device:
+        if not managed and self.transformer.device != hidden_states.device:
             if self.low_vram:
                 # move other transformer to cpu
                 other_tname = (
@@ -285,66 +302,25 @@ class Wan2214bModel(Wan21):
             # we have a hf path, replace it with transformer_2 subfolder
             subfolder_2 = "transformer_2"
 
+        # per-transformer load kwargs; the ARA (if any) applies to the combined
+        # dual model below, offload attaches per transformer after that
+        per_kwargs = self.component_load_kwargs("transformer")
+        per_kwargs["offload"] = 0.0
+        if self.model_config.accuracy_recovery_adapter is not None:
+            per_kwargs["qtype"] = None
+
         self.print_and_status_update("Loading transformer 1")
-        dtype = self.torch_dtype
-        transformer_1 = WanTransformer3DModel.from_pretrained(
-            transformer_path_1,
-            subfolder=subfolder_1,
-            torch_dtype=dtype,
-        ).to(dtype=dtype)
-
+        transformer_1 = WanTransformer3DModel.load(
+            transformer_path_1, subfolder=subfolder_1, **per_kwargs
+        )
         flush()
 
-        if self.model_config.low_vram:
-            # quantize on the device
-            transformer_1.to('cpu', dtype=dtype)
-            flush()
-        else:
-            transformer_1.to(self.device_torch, dtype=dtype)
-            flush()
-
-        if self.model_config.quantize and self.model_config.accuracy_recovery_adapter is None:
-            # todo handle two ARAs
-            self.print_and_status_update("Quantizing Transformer 1")
-            quantize_model(self, transformer_1)
-            flush()
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer 1 to CPU")
-            transformer_1.to("cpu")
-        else:
-            transformer_1.to(self.device_torch)
-
         self.print_and_status_update("Loading transformer 2")
-        dtype = self.torch_dtype
-        transformer_2 = WanTransformer3DModel.from_pretrained(
-            transformer_path_2,
-            subfolder=subfolder_2,
-            torch_dtype=dtype,
-        ).to(dtype=dtype)
-
+        transformer_2 = WanTransformer3DModel.load(
+            transformer_path_2, subfolder=subfolder_2, **per_kwargs
+        )
         flush()
 
-        if self.model_config.low_vram:
-            # quantize on the device
-            transformer_2.to('cpu', dtype=dtype)
-            flush()
-        else:
-            transformer_2.to(self.device_torch, dtype=dtype)
-            flush()
-
-        if self.model_config.quantize and self.model_config.accuracy_recovery_adapter is None:
-            # todo handle two ARAs
-            self.print_and_status_update("Quantizing Transformer 2")
-            quantize_model(self, transformer_2)
-            flush()
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer 2 to CPU")
-            transformer_2.to("cpu")
-        else:
-            transformer_2.to(self.device_torch)
-    
         layer_offloading_transformer = self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0
         # make the combined model
         self.print_and_status_update("Creating DualWanTransformer3DModel")
@@ -365,18 +341,13 @@ class Wan2214bModel(Wan21):
             
         
         if layer_offloading_transformer:
-            MemoryManager.attach(
-                transformer_1,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_transformer_percent,
-                ignore_modules=[transformer_1.scale_shift_table] + [block.scale_shift_table for block in transformer_1.blocks]
-            )
-            MemoryManager.attach(
-                transformer_2,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_transformer_percent,
-                ignore_modules=[transformer_2.scale_shift_table] + [block.scale_shift_table for block in transformer_2.blocks]
-            )
+            for t in (transformer_1, transformer_2):
+                MemoryManager.attach(
+                    t,
+                    self.device_torch,
+                    offload_percent=self.model_config.layer_offloading_transformer_percent,
+                    ignore_modules=t.get_offload_ignore_modules(),
+                )
 
         return transformer
 
@@ -435,20 +406,20 @@ class Wan2214bModel(Wan21):
         return False
 
     def save_model(self, output_path, meta, save_dtype):
+        # comfy-format single-file saves, one per DiT (comfy convention:
+        # separate high/low noise files)
         transformer_combo: DualWanTransformer3DModel = unwrap_model(self.model)
-        transformer_combo.transformer_1.save_pretrained(
-            save_directory=os.path.join(output_path, "transformer"),
-            safe_serialization=True,
+        base = output_path
+        if base.endswith(".safetensors"):
+            base = base[: -len(".safetensors")]
+        metadata = get_meta_for_safetensors(meta, name=self.arch)
+        transformer_combo.transformer_1.save_model(
+            f"{base}_high_noise.safetensors", dtype=save_dtype, metadata=metadata
         )
-        transformer_combo.transformer_2.save_pretrained(
-            save_directory=os.path.join(output_path, "transformer_2"),
-            safe_serialization=True,
+        transformer_combo.transformer_2.save_model(
+            f"{base}_low_noise.safetensors", dtype=save_dtype, metadata=metadata
         )
 
-        meta_path = os.path.join(output_path, "aitk_meta.yaml")
-        with open(meta_path, "w") as f:
-            yaml.dump(meta, f)
-
     def save_lora(
         self,
         state_dict: Dict[str, torch.Tensor],

extensions_built_in/diffusion_models/z_image/z_image.py

@@ -3,7 +3,6 @@ from typing import List, Optional
 
 import huggingface_hub
 import torch
-import yaml
 from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig
 from toolkit.lora_special import LoRASpecialNetwork
 from toolkit.models.base_model import BaseModel
@@ -13,26 +12,19 @@ from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze
-from toolkit.util.quantize import (
-    quantize,
-    get_qtype,
-    quantize_model,
-    dequantize_if_quantized,
-)
 from toolkit.memory_management import MemoryManager
 from toolkit.metadata import get_meta_for_safetensors
-from safetensors.torch import load_file, save_file
+from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
+from toolkit.models.v2.vae.autoencoder_kl import KLVAE
+from safetensors.torch import load_file
 
-from transformers import AutoTokenizer, Qwen3ForCausalLM
-from diffusers import AutoencoderKL
 
 try:
     from diffusers import ZImagePipeline
 
     # our subclass of the diffusers transformer with the universal loading /
-    # quantization mixin (see toolkit/models/classes/_mixin.py)
-    from toolkit.models.v2.z_image import ZImageTransformer2DModel
+    # quantization mixin (see toolkit/models/v2/_mixin.py)
+    from toolkit.models.v2.diffusion_models.z_image import ZImageTransformer2DModel
 except ImportError:
     raise ImportError(
         "Diffusers is out of date. Update diffusers to the latest version by doing pip uninstall diffusers and then pip install -r requirements.txt"
@@ -189,20 +181,13 @@ class ZImageModel(BaseModel):
         # quantization happens inside load_model unless an adapter has to be merged
         # into the full precision weights first (assistant lora / accuracy recovery
         # adapter); those paths quantize after the merge via quantize_model
-        qtype = None
-        if (
-            self.model_config.quantize
-            and self.model_config.assistant_lora_path is None
-            and self.model_config.accuracy_recovery_adapter is None
-        ):
-            qtype = self.model_config.qtype
-
         transformer = ZImageTransformer2DModel.load_model(
             model_path,
             dtype=dtype,
-            qtype=qtype,
-            quantize_device=self.device_torch,
             config_path=base_model_path if self.is_single_file else None,
+            use_comfy_weights=self.model_config.model_kwargs.get(
+                "use_comfy_weights", True
+            ),
         )
         flush()
 
@@ -227,64 +212,18 @@ class ZImageModel(BaseModel):
             if self.model_config.qtype == "qfloat8":
                 self.model_config.qtype = "float8"
 
-        # already quantized inside load_transformer unless an adapter had to merge
-        # into full precision weights first (or the checkpoint was pre-quantized)
-        if self.model_config.quantize and not transformer.aitk_is_quantized:
-            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,
-                ignore_modules=[
-                    transformer.x_pad_token,
-                    transformer.cap_pad_token,
-                ],
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
-
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
 
         self.print_and_status_update("Text Encoder")
-        tokenizer = AutoTokenizer.from_pretrained(
-            base_model_path, subfolder="tokenizer", torch_dtype=dtype
+        tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
+        text_encoder = Qwen3TextEncoder.load(
+            base_model_path, **self.component_load_kwargs("te")
         )
-        text_encoder = Qwen3ForCausalLM.from_pretrained(
-            base_model_path, subfolder="text_encoder", torch_dtype=dtype
-        )
-
-        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,
-            )
-
-        text_encoder.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
 
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKL.from_pretrained(
-            base_model_path, subfolder="vae", torch_dtype=dtype
-        )
+        vae = KLVAE.load_model(base_model_path, dtype=dtype)
 
         self.noise_scheduler = ZImageModel.get_train_scheduler()
 
@@ -427,31 +366,17 @@ class ZImageModel(BaseModel):
         return ZImageTransformer2DModel.get_quantization_exclude_modules()
 
     def save_model(self, output_path, meta, save_dtype):
+        # comfy-format single-file save (the standard save format regardless
+        # of how the model was loaded); prequantized layers keep their
+        # quantized storage
         transformer: ZImageTransformer2DModel = unwrap_model(self.model)
-        if self.is_single_file:
-            # loaded from a single-file checkpoint, save back in that format
-            sd = transformer.state_dict()
-            save_dict = {}
-            for key, value in sd.items():
-                # dequantize any quantized (e.g. torchao) weights so we save plain tensors
-                save_dict[key] = (
-                    dequantize_if_quantized(value).clone().to("cpu", dtype=save_dtype)
-                )
-            save_dict = transformer.convert_state_dict_on_save(save_dict)
-
-            if not output_path.endswith(".safetensors"):
-                output_path += ".safetensors"
-            meta = get_meta_for_safetensors(meta, name=self.arch)
-            save_file(save_dict, output_path, metadata=meta)
-        else:
-            transformer.save_pretrained(
-                save_directory=os.path.join(output_path, "transformer"),
-                safe_serialization=True,
-            )
-
-            meta_path = os.path.join(output_path, "aitk_meta.yaml")
-            with open(meta_path, "w") as f:
-                yaml.dump(meta, f)
+        if not output_path.endswith(".safetensors"):
+            output_path += ".safetensors"
+        transformer.save_model(
+            output_path,
+            dtype=save_dtype,
+            metadata=get_meta_for_safetensors(meta, name=self.arch),
+        )
 
     def get_loss_target(self, *args, **kwargs):
         noise = kwargs.get("noise")
@@ -464,16 +389,5 @@ class ZImageModel(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["layers"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd

extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py

@@ -9,11 +9,8 @@ import torch.nn.functional as F
 import yaml
 from toolkit.basic import flush
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
-from toolkit.memory_management import MemoryManager
 
-from transformers import AutoTokenizer, Qwen3ForCausalLM
+from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
 from toolkit.models.FakeVAE import FakeVAE
 from toolkit.paths import MODELS_PATH
 from safetensors.torch import load_file, save_file
@@ -23,7 +20,7 @@ HF_TOKEN = os.getenv("HF_TOKEN", None)
 
 try:
     from diffusers import ZImagePipeline
-    from diffusers.models.transformers.transformer_z_image import (
+    from toolkit.models.v2.diffusion_models.z_image import (
         ZImageTransformer2DModel as ZImageTransformer2DModelOriginal,
     )
     from diffusers.models.modeling_outputs import Transformer2DModelOutput
@@ -442,57 +439,16 @@ class ZImageL2PModel(ZImageModel):
             if self.model_config.qtype == "qfloat8":
                 self.model_config.qtype = "float8"
 
-        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,
-                ignore_modules=[
-                    transformer.x_pad_token,
-                    transformer.cap_pad_token,
-                ],
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
 
         flush()
 
         self.print_and_status_update("Text Encoder")
-        tokenizer = AutoTokenizer.from_pretrained(
-            base_model_path, subfolder="tokenizer", torch_dtype=dtype
+        tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
+        text_encoder = Qwen3TextEncoder.load(
+            base_model_path, **self.component_load_kwargs("te")
         )
-        text_encoder = Qwen3ForCausalLM.from_pretrained(
-            base_model_path, subfolder="text_encoder", torch_dtype=dtype
-        )
-
-        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,
-            )
-
-        text_encoder.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
 
         self.print_and_status_update("Loading VAE")
         # vae = AutoencoderKL.from_pretrained(

extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py

@@ -13,21 +13,17 @@ from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
-from toolkit.memory_management import MemoryManager
 from safetensors.torch import load_file
 from optimum.quanto import QTensor
 from toolkit.metadata import get_meta_for_safetensors
 from safetensors.torch import load_file, save_file
-from transformers import AutoTokenizer, Qwen3ForCausalLM
+from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
 from diffusers import AutoencoderKL
 from toolkit.models.FakeVAE import FakeVAE
 from .zeta_chroma_transformer import ZImageDCT, ZImageDCTParams, vae_flatten, vae_unflatten, prepare_latent_image_ids, make_text_position_ids
 from .zeta_chroma_pipeline import ZetaChromaPipeline
 
 
-
 scheduler_config = {
     "num_train_timesteps": 1000,
     "use_dynamic_shifting": False,
@@ -94,10 +90,6 @@ class ZetaChromaModel(BaseModel):
 
         transformer_state_dict = load_file(transformer_path, device="cpu")
 
-        # cast to dtype
-        for key in transformer_state_dict:
-            transformer_state_dict[key] = transformer_state_dict[key].to(dtype)
-        
         # Auto-detect use_x0 from checkpoint
         use_x0 = "__x0__" in transformer_state_dict
 
@@ -109,65 +101,21 @@ class ZetaChromaModel(BaseModel):
             use_x0=use_x0,
         )
 
-        with torch.device("meta"):
-            transformer = ZImageDCT(model_params)
-            
-        transformer.load_state_dict(transformer_state_dict, assign=True)
+        transformer = ZImageDCT.load_from_state_dict(
+            transformer_state_dict, dtype, config=model_params
+        )
         del transformer_state_dict
 
-        transformer.to(self.quantize_device, dtype=dtype)
-
-        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,
-                ignore_modules=[
-                    transformer.x_pad_token,
-                    transformer.cap_pad_token,
-                ],
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
 
         flush()
 
         self.print_and_status_update("Text Encoder")
-        tokenizer = AutoTokenizer.from_pretrained(
-            base_model_path, subfolder="tokenizer", torch_dtype=dtype
+        tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
+        text_encoder = Qwen3TextEncoder.load(
+            base_model_path, **self.component_load_kwargs("te")
         )
-        text_encoder = Qwen3ForCausalLM.from_pretrained(
-            base_model_path, subfolder="text_encoder", torch_dtype=dtype
-        )
-
-        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,
-            )
-
-        text_encoder.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
 
         self.print_and_status_update("Loading VAE")
         vae = FakeVAE(scaling_factor=1.0)
@@ -201,8 +149,10 @@ class ZetaChromaModel(BaseModel):
             pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        # just to make sure everything is on the right device and dtype
-        text_encoder[0].to(self.device_torch)
+        # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
+        # to the gpu on demand
+        if not self.low_vram:
+            text_encoder[0].to(self.device_torch)
         text_encoder[0].requires_grad_(False)
         text_encoder[0].eval()
         flush()
@@ -368,16 +318,5 @@ class ZetaChromaModel(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["layers"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
-    def convert_lora_weights_before_load(self, state_dict):
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd

extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py

@@ -6,6 +6,8 @@ from typing import List, Optional
 import math
 
 import torch
+
+from toolkit.models.v2._mixin import OstrisModelMixin
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
@@ -449,7 +451,14 @@ class SimpleMLPAdaLN(nn.Module):
         return self.final_layer(x)
 
 
-class ZImageDCT(nn.Module):
+class ZImageDCT(nn.Module, OstrisModelMixin):
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]
+
+    def get_offload_ignore_modules(self):
+        return [self.x_pad_token, self.cap_pad_token]
+
     def __init__(self, params: ZImageDCTParams):
         super().__init__()
         self.config = FakeConfig()

testing/stage_profiler.py

@@ -0,0 +1,157 @@
+"""Per-stage speed / memory / CPU profiler for the model loading test harness.
+
+A StageProfiler splits a run into named stages (the harness hooks
+BaseModel.print_and_status_update so every holder status line — "Loading
+transformer", "Quantizing (convrot8)", "Loading VAE", ... — opens a new
+stage). For every stage it records:
+
+  - seconds          wall time
+  - vram_peak_gb     torch.cuda.max_memory_allocated during the stage
+  - vram_reserved_gb torch.cuda.max_memory_reserved during the stage
+  - vram_end_gb      memory_allocated at the stage boundary (resident model)
+  - rss_peak_gb      peak process RSS (sampled at 50ms)
+  - cpu_avg_pct /    process CPU utilization (100 = one full core), sampled
+    cpu_max_pct      at 50ms — this is what shows where quantization burns CPU
+  - threads_max      peak thread count
+
+A background sampler thread does the RSS/CPU sampling so short spikes inside
+a stage are caught, not just the boundary values.
+"""
+
+import threading
+import time
+
+import psutil
+import torch
+
+_GB = 1024**3
+
+
+class _Sampler:
+    """Background thread sampling RSS / CPU% / thread count at 50ms."""
+
+    def __init__(self):
+        self.proc = psutil.Process()
+        self._lock = threading.Lock()
+        self._stop = False
+        self.reset()
+        self.proc.cpu_percent(None)  # prime the cpu_percent window
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._thread.start()
+
+    def reset(self):
+        with self._lock:
+            self.rss_peak = self.proc.memory_info().rss
+            self.cpu_samples = []
+            self.threads_max = self.proc.num_threads()
+
+    def _run(self):
+        while not self._stop:
+            try:
+                rss = self.proc.memory_info().rss
+                cpu = self.proc.cpu_percent(None)
+                threads = self.proc.num_threads()
+                with self._lock:
+                    self.rss_peak = max(self.rss_peak, rss)
+                    self.cpu_samples.append(cpu)
+                    self.threads_max = max(self.threads_max, threads)
+            except Exception:
+                pass
+            time.sleep(0.05)
+
+    def stats(self):
+        with self._lock:
+            samples = self.cpu_samples or [0.0]
+            return {
+                "rss_peak_gb": round(self.rss_peak / _GB, 3),
+                "cpu_avg_pct": round(sum(samples) / len(samples), 1),
+                "cpu_max_pct": round(max(samples), 1),
+                "threads_max": self.threads_max,
+            }
+
+    def stop(self):
+        self._stop = True
+
+
+class StageProfiler:
+    def __init__(self, device):
+        self.device = torch.device(device) if torch.cuda.is_available() else None
+        self.stages = []
+        self._current = None
+        self._sampler = _Sampler()
+
+    def _cuda_sync_reset(self):
+        if self.device is not None:
+            torch.cuda.synchronize(self.device)
+            torch.cuda.reset_peak_memory_stats(self.device)
+
+    def stage(self, name: str):
+        """Close the current stage (if any) and open a new one."""
+        self._close()
+        self._cuda_sync_reset()
+        self._sampler.reset()
+        self._current = {"name": str(name), "t0": time.perf_counter()}
+
+    def _close(self):
+        if self._current is None:
+            return
+        cur = self._current
+        self._current = None
+        seconds = time.perf_counter() - cur["t0"]
+        entry = {"name": cur["name"], "seconds": round(seconds, 3)}
+        if self.device is not None:
+            torch.cuda.synchronize(self.device)
+            entry["vram_peak_gb"] = round(
+                torch.cuda.max_memory_allocated(self.device) / _GB, 3
+            )
+            entry["vram_reserved_gb"] = round(
+                torch.cuda.max_memory_reserved(self.device) / _GB, 3
+            )
+            entry["vram_end_gb"] = round(
+                torch.cuda.memory_allocated(self.device) / _GB, 3
+            )
+        entry.update(self._sampler.stats())
+        self.stages.append(entry)
+
+    def finish(self):
+        self._close()
+        self._sampler.stop()
+        return self.stages
+
+
+def profile_top(profile, limit=40):
+    """Top functions from a cProfile.Profile by cumulative time, as dicts."""
+    import pstats
+
+    stats = pstats.Stats(profile)
+    stats.sort_stats("cumulative")
+    rows = []
+    for func in stats.fcn_list[: limit * 3]:
+        cc, nc, tt, ct, _ = stats.stats[func]
+        filename, line, name = func
+        # drop the profiler/exec wrappers and trivial rows
+        if name in ("<module>",) or ct < 0.05:
+            continue
+        # cProfile catches this module's own 50ms sampling thread; its idle
+        # sleep loop and psutil polling otherwise show up as a giant fake
+        # "time.sleep" hotspot spanning the whole stage
+        if "stage_profiler" in filename or "psutil" in filename or "_pslinux" in filename:
+            continue
+        if name == "<built-in method time.sleep>":
+            continue
+        # shorten site-packages / repo paths to keep the report readable
+        for marker in ("site-packages/", "ai-toolkit/"):
+            if marker in filename:
+                filename = filename.split(marker, 1)[1]
+                break
+        rows.append(
+            {
+                "func": f"{filename}:{line}({name})",
+                "ncalls": nc,
+                "tottime": round(tt, 3),
+                "cumtime": round(ct, 3),
+            }
+        )
+        if len(rows) >= limit:
+            break
+    return rows

testing/test_model_loading.py

@@ -0,0 +1,702 @@
+"""Per-arch model loading + inference smoke test (toolkit/models/v2/PLANNING.md).
+
+Loads one registered arch through its normal loading path (the same
+get_model_class -> ModelClass(...).load_model() flow training uses), runs one
+small sample generation, and asserts an output file was produced.
+
+Usage:
+  python testing/test_model_loading.py --arch zimage          # one arch, in-process
+  python testing/test_model_loading.py --all                  # every registered arch,
+                                                              # one subprocess each (full
+                                                              # unload between archs)
+  --allow-download   permit hub downloads (default: HF_HUB_OFFLINE=1, so archs
+                     whose weights are not local/cached report SKIP)
+  --list             list registered archs
+  --device cuda:0
+
+Add a new model type by adding an entry to MODEL_TESTS.
+"""
+
+import argparse
+import glob
+import json
+import os
+import subprocess
+import sys
+import time
+
+TOOLKIT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, TOOLKIT_ROOT)
+
+from dotenv import load_dotenv
+
+# repo .env carries HF_TOKEN / HF_HOME / MODELS_PATH etc., same as run.py
+load_dotenv(os.path.join(TOOLKIT_ROOT, ".env"))
+
+OUTPUT_ROOT = os.path.join(TOOLKIT_ROOT, "testing", ".model_test_outputs")
+
+# arch -> {"model": ModelConfig kwargs, "sample": GenerateImageConfig kwargs}
+# Keep samples tiny: this asserts the load/encode/denoise/decode/save path
+# works, not quality.
+IMG = {"width": 512, "height": 512, "num_inference_steps": 8, "seed": 42}
+VID = {"width": 256, "height": 256, "num_inference_steps": 6, "seed": 42, "num_frames": 9}
+
+MODEL_TESTS = {
+    "zimage": {
+        "model": {"name_or_path": "Tongyi-MAI/Z-Image-Turbo"},
+        "sample": {**IMG, "guidance_scale": 1.0},
+    },
+    "qwen_image": {
+        # 20B: quantize to fit a 32GB card
+        "model": {"name_or_path": "Qwen/Qwen-Image", "quantize": True, "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0},
+    },
+    "krea2": {
+        "model": {"name_or_path": "krea/Krea-2-Turbo", "quantize": True, "quantize_te": True},
+        "sample": {**IMG, "guidance_scale": 1.0},
+    },
+    "boogu_image": {
+        # native ~1024; 512/low-step/high-CFG degenerates to a black frame
+        "model": {"name_or_path": "Boogu/Boogu-Image-0.1-Base", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42},
+    },
+    "ernie_image": {
+        "model": {"name_or_path": "baidu/ERNIE-Image", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42},
+    },
+    "mageflow": {
+        "model": {"name_or_path": "microsoft/Mage-Flow-Base", "quantize": True, "quantize_te": True},
+        "sample": IMG,
+    },
+    "ideogram4": {
+        "model": {"name_or_path": "ideogram-ai/ideogram-4-fp8", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42},
+    },
+    "hidream_o1": {
+        "model": {"name_or_path": "HiDream-ai/HiDream-O1-Image", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42},
+    },
+    "anima": {
+        "model": {"name_or_path": "circlestone-labs/Anima-Base-v1.0-Diffusers"},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.5, "seed": 42},
+    },
+    "wan21": {
+        "model": {"name_or_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"},
+        "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 5.0, "seed": 42, "num_frames": 17},
+    },
+    "wan22_5b": {
+        "model": {"name_or_path": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "quantize": True, "quantize_te": True},
+        "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 5.0, "seed": 42, "num_frames": 17},
+    },
+    "ltx2.3": {
+        # even quantized, the 22B stack does not fit a 32GB card — needs the big GPU
+        "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 3.0, "seed": 42, "num_frames": 25},
+    },
+    # single-file / comfy-layout archs: weights resolve under MODELS_PATH (or
+    # download there with --allow-download)
+    "chroma": {
+        "model": {"name_or_path": "lodestones/Chroma1-HD", "quantize": True, "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 26, "guidance_scale": 4.0},
+    },
+    "flux_kontext": {
+        "model": {"name_or_path": "black-forest-labs/FLUX.1-Kontext-dev", "quantize": True, "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 2.5},
+        "needs_control_image": True,
+    },
+    "flux2_klein_4b": {
+        "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-4B", "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0},
+    },
+    # ---- coverage for every UI-default arch (big downloads on first run) ----
+    "wan22_14b": {
+        "model": {"name_or_path": "ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True},
+        "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 3.5, "seed": 42, "num_frames": 17},
+    },
+    "wan22_14b_i2v": {
+        "model": {"name_or_path": "ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True},
+        "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 3.5, "seed": 42, "num_frames": 17},
+        "needs_control_image": True,
+    },
+    "wan21_i2v": {
+        "model": {"name_or_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", "quantize": True, "quantize_te": True},
+        "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 5.0, "seed": 42, "num_frames": 17},
+        "needs_control_image": True,
+    },
+    "hidream": {
+        "model": {"name_or_path": "HiDream-ai/HiDream-I1-Full", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42},
+    },
+    "hidream_e1": {
+        # editing seq budget fits 768x768 (source+target concat)
+        "model": {"name_or_path": "HiDream-ai/HiDream-E1-1", "quantize": True, "quantize_te": True},
+        "sample": {"width": 768, "height": 768, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42},
+        "needs_control_image": True,
+    },
+    "nucleus_image": {
+        "model": {"name_or_path": "NucleusAI/Nucleus-Image", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42},
+    },
+    "omnigen2": {
+        "model": {"name_or_path": "OmniGen2/OmniGen2", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42},
+    },
+    "ltx2.5": {
+        "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 3.0, "seed": 42, "num_frames": 25},
+    },
+    "flux2": {
+        "model": {"name_or_path": "black-forest-labs/FLUX.2-dev", "quantize": True, "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0},
+    },
+    "flux2_klein_9b": {
+        "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-9B", "quantize": True, "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0},
+    },
+    "prx_pixel": {
+        "model": {"name_or_path": "Photoroom/prxpixel-t2i", "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0},
+    },
+    "zeta_chroma": {
+        "model": {"name_or_path": "lodestones/Zeta-Chroma/zeta-chroma-base-x0-pixel-dino-distance.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize": True, "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0},
+    },
+    "zimage_l2p": {
+        "model": {"name_or_path": "zhen-nan/L2P/model-1k-merge.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize_te": True},
+        "sample": {**IMG, "guidance_scale": 1.0},
+    },
+    "qwen_image_edit": {
+        "model": {"name_or_path": "Qwen/Qwen-Image-Edit", "quantize": True, "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0},
+        "needs_control_image": True,
+    },
+    "qwen_image_edit_plus": {
+        "model": {"name_or_path": "Qwen/Qwen-Image-Edit-2509", "quantize": True, "quantize_te": True},
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0},
+        "needs_control_image": True,
+    },
+    # ---- legacy monolith archs (components adopted into v2 on load) ----
+    "sd1": {
+        "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 7.5, "seed": 42},
+    },
+    "sdxl": {
+        "model": {"name_or_path": "stabilityai/stable-diffusion-xl-base-1.0"},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 6.0, "seed": 42},
+    },
+    "ace_step_15": {
+        "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0, "seed": 42},
+    },
+    "f-lite": {
+        "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True},
+        "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42},
+    },
+}
+
+SKIP_MARKERS = (
+    "couldn't connect",
+    "offline mode",
+    "hf_hub_offline",
+    "cannot find the requested files",
+    "not found in cache",
+    "localentrynotfounderror",
+    "does not appear to have a file named",
+    "404 client error",
+    "entrynotfounderror",
+    "gatedrepoerror",
+    "cannot access gated repo",
+    "repositorynotfounderror",
+)
+
+
+def classify_error(err: BaseException) -> str:
+    text = f"{type(err).__name__}: {err}".lower()
+    if any(m in text for m in SKIP_MARKERS):
+        return "SKIP"
+    if isinstance(err, FileNotFoundError):
+        return "SKIP"
+    return "FAIL"
+
+
+def run_one(arch: str, device: str, allow_download: bool) -> dict:
+    entry = MODEL_TESTS[arch]
+    out_dir = os.path.join(OUTPUT_ROOT, arch.replace("/", "_").replace(":", "_"))
+    os.makedirs(out_dir, exist_ok=True)
+    for old in glob.glob(os.path.join(out_dir, "*")):
+        os.remove(old)
+
+    from toolkit.config_modules import GenerateImageConfig, ModelConfig
+    from toolkit.util.get_model import get_model_class
+
+    # quantization always tests the convrot8 backend (per-entry override wins)
+    model_kwargs = dict(entry["model"])
+    if model_kwargs.get("quantize"):
+        model_kwargs.setdefault("qtype", "convrot8")
+    if model_kwargs.get("quantize_te"):
+        model_kwargs.setdefault("qtype_te", "convrot8")
+
+    model_config = ModelConfig(arch=arch, dtype="bf16", **model_kwargs)
+    ModelClass = get_model_class(model_config)
+    from toolkit.util.get_model import LEGACY_ARCHS
+
+    if (
+        getattr(ModelClass, "arch", None) not in (arch, model_config.arch)
+        and model_config.arch not in LEGACY_ARCHS
+    ):
+        raise ValueError(
+            f"arch {arch!r} resolved to {ModelClass.__name__} "
+            f"(arch={getattr(ModelClass, 'arch', None)!r}) — registry mismatch"
+        )
+
+    sampler = None
+    if hasattr(ModelClass, "get_train_scheduler"):
+        sampler = ModelClass.get_train_scheduler()
+    else:
+        # legacy monolith archs build their scheduler the way training does
+        from toolkit.sampler import get_sampler
+
+        legacy_arch = "sd"
+        if model_config.is_pixart:
+            legacy_arch = "pixart"
+        elif model_config.is_flux:
+            legacy_arch = "flux"
+        elif model_config.is_lumina2:
+            legacy_arch = "lumina2"
+        sampler = get_sampler(
+            "ddpm",
+            {
+                "prediction_type": "v_prediction"
+                if model_config.is_v_pred
+                else "epsilon",
+            },
+            arch=legacy_arch,
+        )
+
+    sd = ModelClass(
+        device=device,
+        model_config=model_config,
+        dtype="bf16",
+        noise_scheduler=sampler,
+    )
+
+    # ---- per-stage speed / VRAM / CPU instrumentation + cProfile ----
+    import cProfile
+
+    from testing.stage_profiler import StageProfiler, profile_top
+
+    prof = StageProfiler(device)
+    if hasattr(sd, "print_and_status_update"):
+        # every holder announces its stages through print_and_status_update;
+        # each status line opens a new profiler stage
+        _orig_status = sd.print_and_status_update
+
+        def _status_hook(msg, *a, **k):
+            prof.stage(str(msg))
+            return _orig_status(msg, *a, **k)
+
+        sd.print_and_status_update = _status_hook
+
+    prof.stage("load: init")
+    load_profile = cProfile.Profile()
+    load_profile.enable()
+    t_load0 = time.perf_counter()
+    sd.load_model()
+    load_seconds = time.perf_counter() - t_load0
+    load_profile.disable()
+    load_profile.dump_stats(os.path.join(out_dir, "load.prof"))
+
+    sample_kwargs = dict(entry["sample"])
+    if entry.get("needs_control_image"):
+        # edit/kontext models require a control image; a flat gray input is fine
+        from PIL import Image
+
+        ctrl_path = os.path.join(out_dir, ".ctrl.png")
+        Image.new(
+            "RGB", (sample_kwargs["width"], sample_kwargs["height"]), (128, 128, 128)
+        ).save(ctrl_path)
+        sample_kwargs["ctrl_img"] = ctrl_path
+    gen = GenerateImageConfig(
+        prompt="a photo of a cat sitting on a wooden table",
+        output_folder=out_dir,
+        # the GenerateImageConfig default for output_ext is the Literal type
+        # alias itself; real callers always pass one
+        output_ext="png",
+        **sample_kwargs,
+    )
+    gen_kwargs = {}
+    if not hasattr(ModelClass, "get_train_scheduler"):
+        # the legacy monolith takes the sampler NAME at generate time
+        gen_kwargs["sampler"] = "ddpm"
+
+    prof.stage("generate")
+    gen_profile = cProfile.Profile()
+    gen_profile.enable()
+    t_gen0 = time.perf_counter()
+    sd.generate_images([gen], **gen_kwargs)
+    gen_seconds = time.perf_counter() - t_gen0
+    gen_profile.disable()
+    gen_profile.dump_stats(os.path.join(out_dir, "gen.prof"))
+
+    # ---- 100% offload forward: attach full layer offloading to the DiT and
+    # text encoder(s), then run a 1-step sample so the offload path (TE
+    # encode + denoise + decode with every layer staged from cpu) is proven
+    # for every arch ----
+    import torch
+
+    from toolkit.memory_management import MemoryManager
+
+    prof.stage("offload 100%: attach")
+    offload_targets = []
+    model_module = getattr(sd, "model", None)
+    if model_module is not None:
+        subs = [
+            getattr(model_module, name, None)
+            for name in ("transformer_1", "transformer_2")
+        ]
+        if all(s is not None for s in subs):
+            # dual-DiT wrappers (wan22 14b): pipelines hold and .to() the
+            # sub-transformers directly, so each needs its own manager
+            offload_targets.extend(subs)
+        else:
+            offload_targets.append(model_module)
+    tes = getattr(sd, "text_encoder", None)
+    for te in tes if isinstance(tes, list) else [tes]:
+        offload_targets.append(te)
+    # auxiliary conditioning stacks where archs have them (ltx connectors,
+    # i2v vision towers)
+    offload_targets.append(getattr(sd, "image_encoder", None))
+    offload_targets.append(getattr(getattr(sd, "pipeline", None), "connectors", None))
+    attached = 0
+    for m in offload_targets:
+        if not isinstance(m, torch.nn.Module):
+            continue
+        if type(m).__name__.startswith("Fake"):
+            continue
+        if next(m.parameters(), None) is None:
+            continue
+        get_ignore = getattr(m, "get_offload_ignore_modules", None)
+        ignore = get_ignore() if callable(get_ignore) else None
+        MemoryManager.attach(
+            m,
+            torch.device(device),
+            offload_percent=1.0,
+            ignore_modules=list(ignore or []),
+        )
+        attached += 1
+
+    prof.stage("offload 100%: generate")
+    offload_profile = cProfile.Profile()
+    offload_profile.enable()
+    t_off0 = time.perf_counter()
+    # 2 steps, not 1: diffusers' shift_terminal stretch divides by
+    # one_minus_z[-1], which is zero for a single step — NaN timesteps and a
+    # NaN image would mask real numerical breakage in the offload path
+    gen_offload = GenerateImageConfig(
+        prompt="a photo of a cat sitting on a wooden table",
+        output_folder=out_dir,
+        output_ext="png",
+        **{**sample_kwargs, "num_inference_steps": 2},
+    )
+    sd.generate_images([gen_offload], **gen_kwargs)
+    offload_seconds = time.perf_counter() - t_off0
+    offload_profile.disable()
+    offload_profile.dump_stats(os.path.join(out_dir, "offload.prof"))
+
+    stages = prof.finish()
+
+    produced = [
+        p
+        for p in glob.glob(os.path.join(out_dir, "*"))
+        if os.path.isfile(p)
+        and os.path.getsize(p) > 1024
+        and not p.endswith((".txt", ".prof", ".json"))
+    ]
+    if not produced:
+        raise RuntimeError(f"no output file produced in {out_dir}")
+
+    result = {
+        "arch": arch,
+        "status": "PASS",
+        "outputs": produced,
+        "qtype": model_config.qtype if model_kwargs.get("quantize") else None,
+        "qtype_te": model_config.qtype_te if model_kwargs.get("quantize_te") else None,
+        "load_seconds": round(load_seconds, 3),
+        "gen_seconds": round(gen_seconds, 3),
+        "offload_seconds": round(offload_seconds, 3),
+        "offload_modules_attached": attached,
+        "stages": stages,
+        "profile": {
+            "load_top": profile_top(load_profile),
+            "gen_top": profile_top(gen_profile, limit=15),
+            "offload_top": profile_top(offload_profile, limit=15),
+            "load_prof": os.path.join(out_dir, "load.prof"),
+            "gen_prof": os.path.join(out_dir, "gen.prof"),
+            "offload_prof": os.path.join(out_dir, "offload.prof"),
+        },
+        "env": {
+            "torch_num_threads": torch.get_num_threads(),
+            "cpu_count": os.cpu_count(),
+            "omp_num_threads": os.environ.get("OMP_NUM_THREADS"),
+        },
+    }
+    with open(os.path.join(out_dir, "metrics.json"), "w") as f:
+        json.dump(result, f, indent=2)
+    return result
+
+
+def format_stage_table(stages) -> str:
+    """Aligned markdown table: stage column left-aligned, numbers right-aligned
+    with fixed decimals so they line up in the terminal too."""
+    headers = [
+        "stage", "seconds", "vram peak GB", "vram resv GB", "vram end GB",
+        "rss peak GB", "cpu avg %", "cpu max %", "threads",
+    ]
+
+    def fmt(value, decimals):
+        if value is None or value == "-":
+            return "-"
+        return f"{value:.{decimals}f}"
+
+    rows = [
+        [
+            s["name"],
+            fmt(s.get("seconds"), 2),
+            fmt(s.get("vram_peak_gb", "-"), 2),
+            fmt(s.get("vram_reserved_gb", "-"), 2),
+            fmt(s.get("vram_end_gb", "-"), 2),
+            fmt(s.get("rss_peak_gb"), 2),
+            fmt(s.get("cpu_avg_pct"), 1),
+            fmt(s.get("cpu_max_pct"), 1),
+            str(s.get("threads_max", "-")),
+        ]
+        for s in stages
+    ]
+
+    widths = [
+        max(len(headers[i]), *(len(r[i]) for r in rows)) if rows else len(headers[i])
+        for i in range(len(headers))
+    ]
+
+    def line(cells):
+        out = []
+        for i, cell in enumerate(cells):
+            # stage name left-aligned, everything else right-aligned
+            out.append(cell.ljust(widths[i]) if i == 0 else cell.rjust(widths[i]))
+        return "| " + " | ".join(out) + " |"
+
+    sep = "|" + "|".join(
+        (":" + "-" * (w + 1)) if i == 0 else ("-" * (w + 1) + ":")
+        for i, w in enumerate(widths)
+    ) + "|"
+    return "\n".join([line(headers), sep] + [line(r) for r in rows])
+
+
+def write_report(results, path):
+    """Aggregated markdown report: per-arch stage tables + cross-arch summary
+    + load-profile hotspots, for diagnosing load/quantization speedups and
+    CPU usage."""
+
+    def _sum(stages, pred):
+        return round(sum(s["seconds"] for s in stages if pred(s)), 1)
+
+    lines = ["# Model loading test report", ""]
+    lines.append(
+        "Per-arch stage timings, peak VRAM (torch allocated/reserved), peak "
+        "process RSS, and CPU utilization (100 = one core, sampled at 50ms). "
+        "Quantization is convrot8 across the board. Full cProfile dumps sit "
+        "next to each arch's outputs (load.prof / gen.prof; inspect with "
+        "`python -m pstats <file>` or snakeviz)."
+    )
+    lines.append("")
+
+    # ---- summary ----
+    lines += ["## Summary", ""]
+    lines += [
+        "| arch | status | load s | quantize s | generate s | vram peak GB | rss peak GB | quant cpu avg % |",
+        "|---|---|---|---|---|---|---|---|",
+    ]
+    for r in results:
+        if r["status"] != "PASS" or "stages" not in r:
+            lines.append(
+                f"| {r['arch']} | {r['status']} | - | - | - | - | - | - | "
+            )
+            continue
+        stages = r["stages"]
+        quant = [s for s in stages if "uantiz" in s["name"]]
+        vram_peak = max((s.get("vram_peak_gb", 0) for s in stages), default=0)
+        rss_peak = max((s.get("rss_peak_gb", 0) for s in stages), default=0)
+        quant_cpu = (
+            round(sum(s["cpu_avg_pct"] for s in quant) / len(quant), 1)
+            if quant
+            else "-"
+        )
+        lines.append(
+            f"| {r['arch']} | PASS | {r['load_seconds']} "
+            f"| {_sum(quant, lambda s: True)} | {r['gen_seconds']} "
+            f"| {vram_peak} | {rss_peak} | {quant_cpu} |"
+        )
+    lines.append("")
+
+    # ---- per-arch detail ----
+    for r in results:
+        if r["status"] != "PASS" or "stages" not in r:
+            lines += [f"## {r['arch']} — {r['status']}", "", r.get("error", ""), ""]
+            continue
+        lines += [f"## {r['arch']}", ""]
+        lines += [
+            f"load {r['load_seconds']}s, generate {r['gen_seconds']}s, "
+            f"qtype {r.get('qtype')}, qtype_te {r.get('qtype_te')}, "
+            f"torch threads {r['env']['torch_num_threads']}/{r['env']['cpu_count']} cores",
+            "",
+        ]
+        lines += [format_stage_table(r["stages"]), ""]
+        top = r.get("profile", {}).get("load_top", [])[:12]
+        if top:
+            lines += ["Load hotspots (cumulative):", "", "```"]
+            for row in top:
+                lines.append(
+                    f"{row['cumtime']:>9.2f}s  tot {row['tottime']:>8.2f}s  "
+                    f"x{row['ncalls']:<8} {row['func']}"
+                )
+            lines += ["```", ""]
+
+    with open(path, "w") as f:
+        f.write("\n".join(lines))
+    return path
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--arch", type=str, default=None)
+    parser.add_argument("--all", action="store_true")
+    parser.add_argument("--list", action="store_true")
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--allow-download", action="store_true")
+    parser.add_argument("--json-result", type=str, default=None)
+    parser.add_argument(
+        "--report-only",
+        action="store_true",
+        help="rebuild report.md from each arch's saved metrics.json, no runs",
+    )
+    args = parser.parse_args()
+
+    if args.list:
+        for arch in MODEL_TESTS:
+            print(arch)
+        return
+
+    if args.report_only:
+        results = []
+        for arch in MODEL_TESTS:
+            out_dir = os.path.join(
+                OUTPUT_ROOT, arch.replace("/", "_").replace(":", "_")
+            )
+            metrics_path = os.path.join(out_dir, "metrics.json")
+            if os.path.exists(metrics_path):
+                with open(metrics_path) as f:
+                    results.append(json.load(f))
+            else:
+                results.append(
+                    {"arch": arch, "status": "SKIP", "error": "no metrics.json saved"}
+                )
+        report_path = write_report(results, os.path.join(OUTPUT_ROOT, "report.md"))
+        print(f"report: {report_path}")
+        return
+
+    os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
+    if not args.allow_download:
+        os.environ.setdefault("HF_HUB_OFFLINE", "1")
+
+    if args.arch is not None:
+        if args.arch not in MODEL_TESTS:
+            raise SystemExit(
+                f"arch {args.arch!r} is not registered; --list shows options"
+            )
+        try:
+            result = run_one(args.arch, args.device, args.allow_download)
+        except BaseException as err:
+            status = classify_error(err)
+            result = {"arch": args.arch, "status": status, "error": f"{type(err).__name__}: {err}"}
+            if status == "FAIL":
+                import traceback
+
+                traceback.print_exc()
+        if args.json_result:
+            with open(args.json_result, "w") as f:
+                json.dump(result, f)
+        if result["status"] == "PASS" and "stages" in result:
+            print()
+            print(format_stage_table(result["stages"]))
+            print()
+        print(f"[{result['status']}] {args.arch}" + (f" — {result.get('error', '')}" if result["status"] != "PASS" else ""))
+        if result["status"] == "FAIL":
+            sys.exit(1)
+        return
+
+    if not args.all:
+        parser.print_help()
+        return
+
+    # --all: one subprocess per arch so every model fully unloads (clean CUDA
+    # teardown) before the next loads
+    results = []
+    total = len(MODEL_TESTS)
+    sweep_t0 = time.perf_counter()
+    for i, arch in enumerate(MODEL_TESTS, start=1):
+        # flush so tail -f shows headers/progress live when stdout is a file
+        print(f"\n===== {arch} ({i}/{total}) =====", flush=True)
+        result_path = os.path.join(OUTPUT_ROOT, f".{arch.replace('/', '_')}.result.json")
+        cmd = [
+            sys.executable,
+            os.path.abspath(__file__),
+            "--arch",
+            arch,
+            "--device",
+            args.device,
+            "--json-result",
+            result_path,
+        ]
+        if args.allow_download:
+            cmd.append("--allow-download")
+        proc = subprocess.run(cmd, cwd=TOOLKIT_ROOT)
+        if os.path.exists(result_path):
+            with open(result_path) as f:
+                results.append(json.load(f))
+            os.remove(result_path)
+        else:
+            results.append(
+                {"arch": arch, "status": "FAIL", "error": f"subprocess died (exit {proc.returncode})"}
+            )
+
+        # running tally so a long sweep always shows how far along it is
+        tally = {"PASS": 0, "FAIL": 0, "SKIP": 0}
+        for r in results:
+            tally[r["status"]] = tally.get(r["status"], 0) + 1
+        elapsed = time.perf_counter() - sweep_t0
+        eta = (elapsed / i) * (total - i)
+        print(
+            f">>> progress: {i}/{total} ({i * 100 // total}%) — "
+            f"{tally['PASS']} pass, {tally['FAIL']} fail, {tally['SKIP']} skip — "
+            f"elapsed {elapsed / 60:.0f}m, eta ~{eta / 60:.0f}m",
+            flush=True,
+        )
+
+    print("\n===== summary =====")
+    counts = {"PASS": 0, "FAIL": 0, "SKIP": 0}
+    for r in results:
+        counts[r["status"]] = counts.get(r["status"], 0) + 1
+        line = f"[{r['status']}] {r['arch']}"
+        if r["status"] != "PASS":
+            line += f" — {r.get('error', '')[:160]}"
+        print(line)
+    print(f"\n{counts['PASS']} passed, {counts['FAIL']} failed, {counts['SKIP']} skipped")
+
+    report_path = write_report(results, os.path.join(OUTPUT_ROOT, "report.md"))
+    print(f"report: {report_path}")
+    if counts["FAIL"]:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

toolkit/dequantize.py

@@ -46,6 +46,10 @@ def hacked_state_dict(self, *args, **kwargs):
 
 # hacks the state dict so we can dequantize before saving
 def patch_dequantization_on_save(model):
+    # idempotent: re-patching would capture the hacked state_dict as
+    # orig_state_dict and recurse on save
+    if getattr(model, "orig_state_dict", None) is not None:
+        return
     model.orig_state_dict = model.state_dict
     model.state_dict = partial(hacked_state_dict, model)
   

toolkit/memory_management/manager.py

@@ -6,6 +6,7 @@ from .manager_modules import (
     LinearLayerMemoryManager,
     ConvLayerMemoryManager,
     OstrisLinearLayerMemoryManager,
+    EmbeddingLayerMemoryManager,
     _DEVICE_STATE,
 )
 import random
@@ -92,12 +93,59 @@ class MemoryManager:
         module._mm_to = module.to
         module.to = module._memory_manager.memory_managed_to
 
-        # add ignore modules to unmanaged list
+        # a fully offloaded module's parameters all live on cpu, which makes
+        # ModelMixin.device (and pipelines deriving their execution device
+        # from it) report "cpu" and plant latents/timesteps there. Report the
+        # compute device instead via an in-place subclass (same class-swap
+        # pattern as OstrisLinear/adopt_component); detach() restores it.
+        try:
+            module._mm_orig_class = module.__class__
+            managed_cls = type(
+                module.__class__.__name__,
+                (module.__class__,),
+                {
+                    "device": property(
+                        lambda self: self._memory_manager.process_device,
+                        lambda self, value: self.__dict__.__setitem__(
+                            "_mm_device_shadow", value
+                        ),
+                    )
+                },
+            )
+            # transformers keys per-class registries (e.g. the hidden-states
+            # capture specs) by str(model.__class__); make the subclass
+            # stringify identically so those lookups still hit
+            managed_cls.__module__ = module.__class__.__module__
+            managed_cls.__qualname__ = module.__class__.__qualname__
+            module.__class__ = managed_cls
+        except TypeError:
+            # exotic class layouts (__slots__ etc.): keep the original class
+            module._mm_orig_class = None
+
+        # add ignore modules to unmanaged list; they must stay RESIDENT on the
+        # compute device (fp32 tables, pad tokens) — a model attached while
+        # parked on cpu would otherwise feed cpu tensors into gpu math
         for im in ignore_modules:
             module._memory_manager.unmanaged_modules.append(im)
+            try:
+                if isinstance(im, torch.Tensor):
+                    im.data = im.data.to(device)
+                elif isinstance(im, torch.nn.Module):
+                    im.to(device)
+            except Exception:
+                pass
 
         # count ignore modules as processed
         modules_processed = [x for x in ignore_modules]
+
+        # weights tied to an embedding (lm_head <-> embed_tokens) must not be
+        # managed: pinning the linear's weight to cpu strands the (unmanaged)
+        # embedding that shares the same tensor
+        embedding_weight_ptrs = {
+            m.weight.data_ptr()
+            for m in module.modules()
+            if isinstance(m, torch.nn.Embedding)
+        }
         # attach to all modules
         for name, sub_module in module.named_modules():
             for child_name, child_module in sub_module.named_modules():
@@ -110,6 +158,12 @@ class MemoryManager:
                         # randomly skip some modules
                         if random.random() > offload_percent:
                             skip = True
+                    if (
+                        not getattr(child_module, "is_ostris_quantized", False)
+                        and isinstance(getattr(child_module, "weight", None), torch.Tensor)
+                        and child_module.weight.data_ptr() in embedding_weight_ptrs
+                    ):
+                        skip = True
                     if skip:
                         module._memory_manager.unmanaged_modules.append(child_module)
                     else:
@@ -158,12 +212,40 @@ class MemoryManager:
                                 )
                             modules_processed.append(ara)
                     modules_processed.append(child_module)
+                elif (
+                    isinstance(child_module, torch.nn.Embedding)
+                    and child_module not in modules_processed
+                    and child_module.weight.numel()
+                    * child_module.weight.element_size()
+                    > 64 * 1024 * 1024
+                ):
+                    # (the cpu gather is autograd-transparent, so a trainable
+                    # embedding still gets grads — they just land on cpu)
+                    # large frozen vocab table: cpu-resident, rows gathered on
+                    # cpu (only the looked-up tokens cross the bus)
+                    EmbeddingLayerMemoryManager.attach(
+                        child_module, module._memory_manager
+                    )
+                    # the cpu move replaced weight.data, so a tied lm_head no
+                    # longer matches the (original-ptr) tied guard and gets
+                    # managed normally — correct: the shared weight now lives
+                    # on cpu, so the linear must bounce it
+                    modules_processed.append(child_module)
                 elif child_module.__class__.__name__ in UNMANAGED_MODULES or any(
                     inc in child_module.__class__.__name__
                     for inc in UNMANAGED_MODULES_INCLUDES
                 ):
-                    # unmanaged
-                    module._memory_manager.unmanaged_modules.append(child_module)
+                    # unmanaged — but never re-list a module the nested walk
+                    # already managed (a managed Embedding landing here would
+                    # get hauled back to the gpu by memory_managed_to), and
+                    # don't append duplicates
+                    if (
+                        child_module not in modules_processed
+                        and not hasattr(child_module, "_layer_memory_manager")
+                        and child_module
+                        not in module._memory_manager.unmanaged_modules
+                    ):
+                        module._memory_manager.unmanaged_modules.append(child_module)
                 else:
                     continue
 
@@ -179,6 +261,10 @@ class MemoryManager:
         if not hasattr(module, "_memory_manager"):
             return
 
+        if getattr(module, "_mm_orig_class", None) is not None:
+            module.__class__ = module._mm_orig_class
+            del module._mm_orig_class
+
         for unmanaged in module._memory_manager.unmanaged_modules:
             try:
                 if isinstance(unmanaged, torch.Tensor):

toolkit/memory_management/manager_modules.py

@@ -312,6 +312,12 @@ class _BouncingLinearFn(torch.autograd.Function):
             ctx.device = torch.device("cpu")
             return out.to(x.device)
 
+        if x.device != device:
+            # a pipeline that derives its execution device from the module's
+            # (fully offloaded -> "cpu") parameters can feed cpu activations;
+            # compute happens on the staged device and the output stays there
+            x = x.to(device, non_blocking=True)
+
         state = _get_device_state(device)
         # the guard makes current_stream() (used by the staging helpers' event
         # waits/records) resolve to the process device; without it they hit
@@ -461,6 +467,11 @@ class _BouncingConv2dFn(torch.autograd.Function):
             ctx.meta = ("cpu", stride, padding, dilation, groups, target_dtype)
             return out.to(x.device)
 
+        if x.device != device:
+            # cpu activations from a fully-offloaded pipeline: see
+            # _BouncingLinearFn.forward
+            x = x.to(device, non_blocking=True)
+
         state = _get_device_state(device)
         # device guard: see _BouncingLinearFn.forward
         with torch.cuda.device(device):
@@ -631,6 +642,42 @@ class BaseLayerMemoryManager:
             param._is_memory_managed = True
 
 
+class EmbeddingLayerMemoryManager(BaseLayerMemoryManager):
+    """Offloads a (large, frozen) nn.Embedding by keeping the weight on cpu
+    and doing the row gather THERE: instead of staging a multi-GB vocab table
+    to the gpu, only the looked-up rows (tokens x dim, ~KBs) cross the bus.
+    Lookups happen once per prompt, so the cpu gather is free."""
+
+    def __init__(self, module: nn.Module, manager: "MemoryManager"):
+        super().__init__(module, manager)
+
+        # cpu-resident weight; no pinning — the weight never crosses the bus
+        module.weight.data = module.weight.data.to("cpu")
+
+        self._original_forward = module.forward
+
+        def _mm_forward(input_ids, *args, **kwargs):
+            if args or kwargs:
+                return self._original_forward(input_ids, *args, **kwargs)
+            out_device = (
+                input_ids.device
+                if input_ids.device.type == "cuda"
+                else self.manager.process_device
+            )
+            out = F.embedding(
+                input_ids.to("cpu"),
+                self.module.weight,
+                self.module.padding_idx,
+                self.module.max_norm,
+                self.module.norm_type,
+                self.module.scale_grad_by_freq,
+                self.module.sparse,
+            )
+            return out.to(out_device, non_blocking=True)
+
+        module.forward = _mm_forward
+
+
 class LinearLayerMemoryManager(BaseLayerMemoryManager):
     def __init__(
         self,
@@ -739,6 +786,11 @@ class OstrisLinearLayerMemoryManager(BaseLayerMemoryManager):
                 # already resident on device
                 return self._original_forward(x)
 
+            if x.device != device:
+                # cpu activations from a fully-offloaded pipeline: see
+                # _BouncingLinearFn.forward
+                x = x.to(device, non_blocking=True)
+
             state = _get_device_state(device)
             d = state["depth"]
             idx = state["forward_clk"]

toolkit/models/base_model.py

@@ -98,6 +98,9 @@ UNET_IN_CHANNELS = 4  # Stable Diffusion の in_channels は 4 で固定。XLも
 class BaseModel:
     # override these in child classes
     arch = None
+    # rename LoRA keys transformer. <-> diffusion_model. (the ComfyUI-standard
+    # prefix) on save/load
+    lora_keys_use_comfy_prefix = False
 
     def __init__(
             self,
@@ -1624,13 +1627,59 @@ class BaseModel:
                 encoder.to(*args, **kwargs)
         else:
             self.text_encoder.to(*args, **kwargs)
-    
+
+    def component_load_kwargs(self, role: str = "transformer", dtype=None):
+        """kwargs for a v2 module's .load()/.aitk_post_load(), derived from
+        model_config: qtype (with the accuracy recovery adapter recombined),
+        offload fraction, devices, low_vram placement. Roles: "transformer",
+        "te", "vae"."""
+        mc = self.model_config
+        qtype, offload = None, 0.0
+        if role == "transformer":
+            if mc.quantize:
+                qtype = mc.qtype
+                if mc.accuracy_recovery_adapter and "|" not in (qtype or ""):
+                    qtype = f"{qtype}|{mc.accuracy_recovery_adapter}"
+            if mc.layer_offloading:
+                offload = mc.layer_offloading_transformer_percent
+        elif role == "te":
+            if mc.quantize_te:
+                qtype = mc.qtype_te
+            if mc.layer_offloading:
+                offload = mc.layer_offloading_text_encoder_percent
+        if dtype is None:
+            dtype = self.vae_torch_dtype if role == "vae" else self.torch_dtype
+        device = self.te_device_torch if role == "te" else self.device_torch
+        if mc.low_vram and role in ("transformer", "te"):
+            device = "cpu"
+        elif role == "vae":
+            device = self.vae_device_torch
+        return dict(
+            qtype=qtype,
+            offload=offload,
+            dtype=dtype,
+            device=device,
+            quantize_device=self.device_torch,
+            base_model=self,
+            use_comfy_weights=mc.model_kwargs.get("use_comfy_weights", True),
+        )
+
     def convert_lora_weights_before_save(self, state_dict):
         # can be overridden in child classes to convert weights before saving
+        if self.lora_keys_use_comfy_prefix:
+            return {
+                k.replace("transformer.", "diffusion_model."): v
+                for k, v in state_dict.items()
+            }
         return state_dict
-    
+
     def convert_lora_weights_before_load(self, state_dict):
         # can be overridden in child classes to convert weights before loading
+        if self.lora_keys_use_comfy_prefix:
+            return {
+                k.replace("diffusion_model.", "transformer."): v
+                for k, v in state_dict.items()
+            }
         return state_dict
     
     def condition_noisy_latents(self, latents: torch.Tensor, batch:'DataLoaderBatchDTO'):

toolkit/models/loaders/umt5.py

@@ -1,46 +1,29 @@
 from typing import List
-import torch
-from transformers import T5Tokenizer, UMT5EncoderModel
 
-class PatchedT5Tokenizer(T5Tokenizer):
-    def __init__(
-        self,
-        vocab: str | list[tuple[str, float]] | None = None,
-        eos_token="</s>",
-        unk_token="<unk>",
-        pad_token="<pad>",
-        _spm_precompiled_charsmap=None,
-        extra_ids=100,
-        additional_special_tokens=None,
-        **kwargs,
-    ):
-        super().__init__(
-            vocab=vocab,
-            eos_token=eos_token,
-            unk_token=unk_token,
-            pad_token=pad_token,
-            _spm_precompiled_charsmap=None, # this is passing a empty byte string for some reason now.
-            extra_ids=extra_ids,
-            additional_special_tokens=additional_special_tokens,
-            **kwargs,
-        )
+import torch
+
+from toolkit.models.v2.text_encoders.umt5 import (
+    PatchedT5Tokenizer,
+    UMT5TextEncoder,
+)
+
 
 def get_umt5_encoder(
     model_path: str,
     tokenizer_subfolder: str = None,
     encoder_subfolder: str = None,
     torch_dtype: str = torch.bfloat16,
-    comfy_files: List[str] = [
-        "text_encoders/umt5_xxl_fp16.safetensors",
-        "text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors",
-    ],
-) -> UMT5EncoderModel:
-    """
-    Load the UMT5 encoder model from the specified path.
-    """
-    tokenizer = PatchedT5Tokenizer.from_pretrained(model_path, subfolder=tokenizer_subfolder)
+    # reserved for the comfy-weights flip (Phase 2); accepted for
+    # signature compatibility, not consulted yet
+    comfy_files: List[str] = None,
+):
+    """Load the UMT5 tokenizer + encoder. Thin compatibility wrapper around
+    toolkit/models/v2/text_encoders/umt5.py."""
+    tokenizer = UMT5TextEncoder.load_tokenizer(
+        model_path, subfolder=tokenizer_subfolder or ""
+    )
     print(f"Using {model_path} for UMT5 encoder.")
-    text_encoder = UMT5EncoderModel.from_pretrained(
-        model_path, subfolder=encoder_subfolder, torch_dtype=torch_dtype
+    text_encoder = UMT5TextEncoder.load_model(
+        model_path, dtype=torch_dtype, subfolder=encoder_subfolder or ""
     )
     return tokenizer, text_encoder

toolkit/models/v2/PLANNING.md

@@ -0,0 +1,626 @@
+# v2 Model Module Restructure — Planning
+
+## Goal
+
+Every component the toolkit loads (DiTs/transformers, unets, text encoders, vision
+encoders, VAEs, audio VAEs) becomes a class extending one base module in
+`toolkit/models/v2`. `BaseModel` (`toolkit/models/base_model.py`) stays as the
+multimodal holder that each arch in `extensions_built_in/diffusion_models` extends —
+that layer is good. The layer below it is what gets unified: one loading entry point,
+one quantization path, one save path, shared component definitions instead of
+per-model-folder copies.
+
+End state this enables:
+
+- **Live server with model hot-swap**: a resident process where, when a generation or
+  training run requests a different model, the unused components are dropped and the
+  new ones loaded. Shared component classes (same TE/VAE reused across archs) make
+  component-level reuse possible instead of full teardown/reload.
+- **Model loading test suite**: a test that loads each registered arch one at a time
+  and runs an inference pass. Every model type gets added to this suite as it is
+  migrated (see Testing below).
+- **Comfy-aligned weights**: weights live in the ComfyUI folder layout under
+  `MODELS_PATH` (shareable with a ComfyUI install), download there when missing, and
+  saves are comfy-format. Eventually defaults move to comfy / our own prequantized
+  releases for everything.
+
+## Current state (survey 2026-08-27)
+
+Three generations of loading conventions coexist:
+
+1. Legacy monolith `toolkit/stable_diffusion_model.py` (`is_flux` / `is_v3` branches);
+   still the silent fallback in `toolkit/util/get_model.py` when an arch string
+   doesn't match.
+2. `BaseModel` subclass per arch (35 registered classes), each with a hand-written
+   `load_model()` / `save_model()`.
+3. `toolkit/models/v2/_mixin.py` (`OstrisModelMixin`) — the intended fix, currently
+   used by one model (`v2/z_image.py` → z_image extension).
+
+### Duplication highlights
+
+- BFL KL autoencoder: full copies in `flux2/src/autoencoder.py` and
+  `ideogram4/src/vae.py` (header says "Flux2 KL autoencoder").
+- Qwen3-VL text encoder loaded independently in qwen_image, nucleus_image, krea2,
+  ideogram4, mageflow, minimax_h3 — from several different repo sources; krea2
+  hand-patches the vision tower locally.
+- `Qwen3ForCausalLM` TE load: 3 verbatim line-for-line copies
+  (`z_image/z_image.py:257`, `z_image/z_image_l2p_model.py:471`,
+  `zeta_chroma/zeta_chroma_model.py:146`).
+- Flux1 VAE + T5 + CLIP trio loaded 4x from 2 different repos (chroma ×2,
+  flux_kontext, legacy SD path).
+- Comfy-file resolver copy-pasted: `minimax_h3/minimax_h3.py:248`
+  (`_resolve_comfy_file`) → `ltx2/ltx2.py:1254` ("mirrors MinimaxH3Model").
+- `AutoencoderKLQwenImage` latents mean/std handling triplicated (qwen_image,
+  nucleus_image, krea2).
+- `transformer.` ↔ `diffusion_model.` LoRA key rename copy-pasted ~15x in
+  `convert_lora_weights_before_save/load` overrides.
+- Fake CLIP/TE/config stubs redefined in ~5 places
+  (canonical: `toolkit/models/FakeVAE.py`, `toolkit/unloader.py`).
+
+### Inconsistency highlights
+
+- **Quantization, 5 paths**: `quantize_model()` (block-streaming, ARA-aware — ~21
+  users), raw `quantize()` (~10 users, no block streaming/excludes, but the only path
+  honoring `quantize_kwargs`), hidream's hand-rolled block loop, the v2 mixin's own
+  `quantize_`, and bare TE quantization everywhere.
+- **Known bugs**: ~9 sites quantize the TE with `qtype` instead of `qtype_te`
+  (chroma ×2, flux2, flux_kontext, cogview4, wan21, legacy SD, ...);
+  `toolkit/models/loaders/umt5.py` accepts a `comfy_files` param it never uses, so
+  wan21's comfy-TE path is a silent no-op.
+- **Saving, 4 incompatible styles**: diffusers `save_pretrained` folders, flat
+  safetensors, safetensors-inside-diffusers-folder hybrids, and z_image's
+  loaded-format-dependent branch. Dequant-on-save done 3 ways; the
+  `isinstance(v, QTensor)` variant (chroma, flux2, boogu_image, ideogram4) misses
+  torchao and Ostris weights entirely. Every `save_pretrained` override ignores its
+  `save_dtype` argument.
+- Registry: linear scan in `toolkit/util/get_model.py`, silent SD1 fallback on a
+  typo'd arch, eager import of every model file at startup. Second unsynchronized
+  registry in `ui/src/app/jobs/new/options.tsx`.
+
+## Decisions (locked in)
+
+1. **Save format = ComfyUI format.** Single-file safetensors in comfy key layout.
+   Must support saving quantized — primarily convrot8 and nvfp4 (comfy_quant marker
+   format, see `toolkit/util/comfy_quant_import.py`) — and plain bf16, all in comfy
+   format. Loading stays backwards compatible: diffusers dirs, transformers repos,
+   and single files all still digest through `load_model`; only saving standardizes
+   on comfy.
+2. **v2 folder layout mirrors the comfy save path structure**:
+
+   ```
+   toolkit/models/v2/
+     _mixin.py            # base module (OstrisModelMixin, evolving)
+     resolver.py          # comfy-layout weight resolution (lift from minimax_h3)
+     diffusion_models/    # one file per DiT/unet family
+     text_encoders/       # qwen3_vl.py, qwen3.py, t5.py, clip.py, gemma.py, ...
+     vae/                 # flux_kl.py, qwen_image.py, wan.py, audio VAEs, ...
+     vision_encoders/
+   ```
+
+3. **Method names win from `BaseModel`**: `get_transformer_block_names` and
+   `get_quantization_exclude_modules`. The mixin's `get_quantization_block_names`
+   gets renamed to match; resolve the classmethod-vs-instance-method mismatch while
+   doing so.
+4. **Loading policy**:
+   - Per-model special handling is allowed via the hook methods.
+   - If `name_or_path` is a diffusers/transformers source, load it with
+     diffusers/transformers for now. **Step 1 is migrating every model to the v2
+     module format and loader without breaking anything** — same weights, same
+     sources, same results.
+   - Each model declares a `comfy_weight_names` dict keyed per standard
+     `name_or_path`. If the user points at a local folder or a non-standard repo,
+     load it as-is. If `name_or_path` is the standard repo and we have matching
+     comfy weight names, load those instead when any of them exist (locally under
+     `MODELS_PATH` in comfy layout, or downloadable to there).
+   - Eventually the default flips to comfy weights / our own prequantized releases
+     for everything.
+
+## Base module: what `OstrisModelMixin` still needs
+
+The mixin already handles: diffusers dir / hub repo / local single file /
+`org/repo/file.safetensors`, key-conversion hooks on load and save, overridable
+backend hooks for transformers-lib models, block-wise quantize.
+
+To add:
+
+- [x] **Comfy weight spec + resolver.** `aitk_comfy_repo` / `aitk_comfy_weight_names`
+      class attrs + `find_comfy_weights` (local-only until Phase 2); resolution chain
+      generalized from `minimax_h3._resolve_comfy_file` into `v2/resolver.py`:
+      explicit override → `MODELS_PATH` at the repo-relative comfy path → flat at
+      root → recursive walk of the category folder → hub download **to the
+      repo-relative path** (folder stays shareable with ComfyUI, no duplicate
+      downloads).
+- [x] **Automatic prequantized import.** Single-file path sniffs `comfy_quant`
+      markers and routes through `import_comfy_quantized_layers` before
+      `load_state_dict`, including the OstrisLinear missing-key whitelist that
+      minimax_h3 and ltx2 each hand-rolled.
+- [x] **One save path.** `save_model(path, dtype)`: dequantize via
+      `dequantize_if_quantized` (honors dtype), run `convert_state_dict_on_save`,
+      write single-file comfy-layout safetensors. (Quantized-storage saves —
+      convrot8 / nvfp4 with comfy_quant markers — land with Phase 2; diffusers-folder
+      save as an explicit flag still to add.)
+- [x] **Tokenizer/processor declaration** for text encoders
+      (`aitk_tokenizer_repo`/`aitk_processor_repo` + `load_tokenizer`/`load_processor`).
+- [x] Rename quantization hooks to the `BaseModel` spellings (decision 3):
+      `get_transformer_block_names` (classmethod on the module).
+
+## Migration steps
+
+Track progress here; check items off as they land.
+
+### Phase 0 — foundation (done 2026-08-27)
+- [x] Evolve `_mixin.py` per the list above (comfy spec local-only until Phase 2)
+- [x] Create `v2/diffusion_models/`, `v2/text_encoders/`, `v2/vae/`,
+      `v2/vision_encoders/`; move `v2/z_image.py` → `v2/diffusion_models/z_image.py`
+- [x] Lift the comfy resolver out of minimax_h3 into `v2/resolver.py`; point
+      minimax_h3 and ltx2 at it (delete their copies)
+- [x] `BaseModel` default `convert_lora_weights_before_save/load` doing the
+      `transformer.` ↔ `diffusion_model.` rename, gated on the class attr
+      `lora_keys_use_comfy_prefix` (default False, so passthrough models keep
+      their behavior); the ~18 identical overrides replaced with the flag.
+      Custom conversions (anima, hidream_o1, ltx2, wan21) keep their overrides;
+      ltx2's now composes with the flag via super().
+
+### Phase 1 — migrate all models to v2 modules, no behavior change
+Every arch's components become v2 classes; if `name_or_path` is diffusers, it still
+loads via diffusers. Nothing about sources or outputs changes yet. Suggested order
+(worst duplication first), each including its loading test (see Testing):
+
+- [x] `text_encoders/qwen3.py` — Qwen3TextEncoder + `OstrisTransformersMixin`
+      backend + `BaseModel.prepare_text_encoder` policy helper; the 3 verbatim
+      TE stanzas (z_image, z_image_l2p, zeta_chroma) replaced. Verified with
+      real Z-Image weights (load + encode on GPU).
+- [x] `text_encoders/qwen3_vl.py` — Qwen3VLTextEncoder with
+      `drop_vision_tower` / `patch_vision_patch_embed`; the 4 identical
+      `patch_qwen_vl_patch_embed` copies (krea2, mageflow, boogu_image,
+      Qwen3VLCaptioner) consolidated; TE loads migrated in krea2, mageflow,
+      nucleus_image. Still on their own paths: ideogram4 (loads via AutoModel),
+      minimax_h3 (custom truncated/prequantized comfy load — port later),
+      qwen_image (Qwen2.5-VL, needs its own class)
+- [x] `text_encoders/t5.py`, `text_encoders/clip.py` — T5TextEncoder,
+      CLIPTextEncoder, CLIPTextEncoderWithProjection; migrated chroma ×2,
+      flux_kontext, f_light (T5 stanzas → `prepare_text_encoder`, fixing their
+      `qtype` → `qtype_te` bug) and hidream (CLIP ×2 + T5 with subfolder
+      overrides; slow-tokenizer classes preserved via `use_fast=False`)
+- [x] `vae/qwen_image.py` — QwenImageVAE + QwenImageVAEHolderMixin (frame-dim +
+      latents mean/std handling built in, tiling opt-in via
+      `vae_decode_tiled_on_low_vram`); the triplicated encode/decode deleted
+      from qwen_image, nucleus_image, krea2 and all three VAE loads routed
+      through the v2 loader
+- [x] `vae/autoencoder_kl.py` — KLVAE (diffusers AutoencoderKL through the
+      universal loader); migrated the scattered loads in chroma, flux_kontext,
+      f_light, hidream, z_image
+- [x] `vae/flux2_kl.py` — the BFL-style Flux2 KL autoencoder unified from the
+      flux2 + ideogram4 copies (both files deleted; flux2's
+      encode/decode/small-decoder superset + ideogram4's diffusers key
+      converter). Verified bit-identical to both originals (weights, encode/
+      decode outputs, converter mapping) and round-tripped real ae.safetensors
+      weights on GPU. Packing/normalization stays per-model — flux2 packs
+      `(c pi pj)` with BatchNorm running stats, ideogram4 packs `(ph pw c)`
+      with its latent_norm tables; the conventions are incompatible.
+- [x] z_image — transformer, TE (qwen3), and VAE (KLVAE) all on v2 modules.
+      z_image_l2p still has its local progressive-transformer subclass
+      (rebasing it onto the v2 class deferred; its TE is migrated)
+- [x] qwen_image family — `v2/diffusion_models/qwen_image.py` (single-file
+      loads stay on diffusers' from_single_file until the comfy flip) +
+      `v2/text_encoders/qwen25_vl.py` (slow tokenizer preserved); edit
+      variants inherit
+- [x] nucleus_image — `v2/diffusion_models/nucleus_image.py`, TE stanza
+      collapsed to prepare_text_encoder
+- [x] krea2, ideogram4, mageflow — DiTs on the mixin via the `config=`
+      passthrough (holder builds config from model_kwargs/holder state, mixin
+      does build/markers/whitelist/casting; plain nn.Module classes now work
+      with the default builder). krea2 + ideogram4 verified by harness;
+      mageflow untestable while its repo 404s
+- [x] chroma, chroma_radiance — both vendored Chroma classes now carry
+      `OstrisModelMixin` with the block-count sniff moved into a new
+      `aitk_config_from_state_dict` hook (mixin now supports checkpoint-derived
+      configs + `load_from_state_dict` for non-safetensors sources, used by
+      radiance's .pth path). zeta_chroma transformer left as-is: its config
+      depends on holder state (patch_size), not the checkpoint
+- [x] flux_kontext — `v2/diffusion_models/flux.py` (FluxTransformer2DModel);
+      whole model now loads through v2 (transformer, T5, CLIP, KLVAE)
+- [x] flux2 + zeta_chroma DiTs on the mixin via `config=` passthrough
+      (flux2_klein_4b verified by harness). Every arch's DiT now loads
+      through the mixin except: ltx2 family (one-file→two-modules split,
+      helper delegated), anima (diffusers modular pipeline), ace_step
+      (bundled single-file loader), z_image_l2p's local subclass, and the
+      grandfathered legacy stable_diffusion_model archs
+- [x] minimax_h3 (+ ref2va) transformer ported to the mixin: config sniffed
+      from the checkpoint via aitk_config_from_state_dict (adaln_t_table),
+      marker attach + stored-precision load via the new
+      `aitk_cast_on_load = False` knob. Verified on the real pruned convrot
+      file: 200 ConvRot linears, pruned table detected, fp32/fp16/bf16 mix
+      preserved, no meta leftovers. Its TE stays custom (50-layer truncation
+      + key_map). ltx2.5's `_load_quantized_module` now delegates to the
+      mixin's whitelist/meta helper (~30 lines deleted); its full port is
+      blocked on the one-comfy-file → transformer+connectors split, which
+      doesn't fit the per-class single-file shape — revisit with the live
+      server's component model
+- [x] wan21 / wan22 family — `v2/diffusion_models/wan.py`
+      (WanTransformer3DModel, both wan22 dual loads included) +
+      `v2/text_encoders/umt5.py` (UMT5TextEncoder + PatchedT5Tokenizer;
+      `loaders/umt5.py` is now a thin compat shim, `comfy_files` still
+      reserved for Phase 2 — no local comfy umt5 file to verify the key
+      conversion against). wan21's TE `qtype` → `qtype_te` bug fixed via
+      prepare_text_encoder
+- [x] hidream family — vendored transformer carries the mixin;
+      `v2/diffusion_models/hidream.py` wraps the diffusers class for
+      hidream_e1; both load via the switchable `hidream_transformer_class`
+      through `load_model`
+- [x] omnigen2 — vendored transformer carries the mixin, load migrated
+- [x] boogu_image, ernie_image, prx_pixel_t2i — their vendored diffusers-style
+      DiT classes now carry OstrisModelMixin (subfolder + block names on the
+      class) and the holders load via `load_model`
+- [x] f_light — DiT class carries the mixin (`aitk_subfolder="dit_model"`),
+      load migrated
+- [ ] anima — loads through diffusers modular pipelines (AnimaModularPipeline);
+      not a mixin fit, revisit at Phase 2
+- [ ] flux2 DiT — holder-config params classes (Flux2/Klein variants), defer
+      like krea2/mageflow
+- [ ] ace_step — one bundled safetensors holds model+TE+VAE+tokenizer via its
+      own load_models; decomposing into v2 components is its own task
+- [ ] Per-model fixes folded in as each migrates: `qtype_te` bug, dequant-on-save
+      (`dequantize_if_quantized` everywhere), raw-`quantize()` → `quantize_model()`
+
+### Phase 2 — comfy weights become the preferred source
+
+Decisions:
+- Comfy weights come from the Comfy-Org hub repos (per-model repos, comfy
+  layout nested under `split_files/` — stripped when placing files into
+  MODELS_PATH). Repos ship several precision variants of each component.
+- **Selection preference: convrot8 > float8 mixed > float8 > bf16 > fp16**
+  (`resolver.comfy_precision_rank`; nvfp4/unmarked rank last and are only
+  used when explicitly listed). **Local-first**: the best-ranked LOCAL
+  candidate wins; only when no candidate is local is the best-ranked one
+  downloaded.
+- Per-model candidate lists (`aitk_comfy_weight_names`) hold only the
+  variants the class can actually digest. Constraint discovered: comfy
+  convrot/quantized files carry markers on the ORIGINAL module layout (e.g.
+  z_image's fused attention.qkv) — diffusers-layout classes with split
+  modules can't attach them until they grow fused-layout support; until then
+  those models list bf16/fp8 variants only. (Vendored comfy-layout classes —
+  minimax/ltx pattern — take convrot directly.)
+- The standard repo still supplies the config; local dirs and unregistered
+  repos load as-is; `model_kwargs.use_comfy_weights: false` opts out.
+
+- [x] Mechanism: `resolver.comfy_precision_rank` / `comfy_local_rel` /
+      `resolve_comfy_candidates` + `OstrisModelMixin.resolve_comfy_weights`,
+      integrated into `load_model` (comfy file preferred for registered
+      standard repos, downloaded into the shared comfy layout)
+- [x] First wired model: z_image (Comfy-Org/z_image_turbo, bf16 candidate).
+      Verified end-to-end against the real shared ComfyUI folder
+      (MODELS_PATH=/mnt/Models/comfy_models): standard-repo name_or_path
+      resolved to the locally-present comfy file and produced the identical
+      generation to the diffusers-shards load
+- [x] qwen_image wired: its comfy files use the diffusers key layout
+      directly — `fp8mixed` (float8-mixed, rank 1) attaches its 839
+      `float8_e4m3fn` markers straight onto the class; candidates fp8mixed →
+      fp8_e4m3fn (raw cast) → bf16. Verified with real weights: loaded the
+      shared folder's local fp8 file (local-first, no download) and generated
+      correctly. Holder skips re-quantization for prequantized checkpoints.
+- [x] New `float8_e4m3fn` Ostris backend (toolkit/util/float8_quant.py):
+      ComfyUI's fp8 + per-tensor-scale storage with dequantized matmul, in
+      get_ostris_quantizer + comfy import/export. Round-trip verified.
+- [x] wan family wired: comfy wan files (original key layout) convert via
+      diffusers' own `convert_wan_transformer_to_diffusers` (rename-only, so
+      quantized weight/scale keys ride along with their modules). Candidate
+      keys support `(repo, subfolder)` tuples for wan2.2 A14B's dual DiTs
+      (transformer = high noise, transformer_2 = low noise) and per-entry
+      comfy-repo overrides ({"repo": ..., "files": [...]}) since wan2.1 and
+      2.2 files live in different Comfy-Org repos. Wired: 2.2 TI2V-5B,
+      T2V/I2V-A14B (fp8_scaled), 2.1 T2V 1.3B/14B, I2V 480P/720P. Verified
+      with real weights: wan21 1.3B (downloaded comfy bf16) and wan22 5B
+      (local comfy fp16) both load through the converter and generate video.
+      Fix along the way: the mixin's meta build now uses accelerate
+      init_empty_weights (params meta, buffers real) so init-computed
+      non-persistent buffers like wan's rope tables materialize.
+- [x] Legacy ComfyUI scaled-fp8 support (`scaled_fp8` marker + per-layer
+      fp8 weight / scalar scale_weight, e.g. every wan *_fp8_scaled file):
+      imports onto the float8 backend; scale_input (activation quant) is
+      dropped, matmuls run dequantized.
+- [x] wan comfy-format saves: `convert_state_dict_on_save` inverts diffusers'
+      rename table (base/t2v/i2v; vace/animate excluded — their reverse
+      mappings collide). Round-trip verified on both real comfy files
+      (exact; the 2.1 file's legacy model.diffusion_model. prefix drops per
+      the modern convention) and with real weights (load → save → 825-key
+      original-layout file → reload bit-equal). wan21 + wan22_5b save one
+      comfy file; wan22_14b saves the comfy-standard _high_noise/_low_noise
+      pair instead of two diffusers folders.
+- [ ] Wire remaining archs' candidate lists (chroma/others as their key
+      conversions are verified per file)
+- [x] Fused-layout quantized attach for diffusers-split classes:
+      `split_fused_quantized_keys` / `fuse_split_quantized_keys`
+      (comfy_quant_import) do exact out-dim row surgery on quantized comfy
+      entries for all three formats (int8 rows+scales slice; fp8 scalar and
+      nvfp4 per-tensor scales shared; nvfp4 block scales
+      unswizzle→split→reswizzle). z_image's load/save converters use them, so
+      its convrot8 candidate is live and top-ranked. Unit-verified exact both
+      directions.
+- [x] Comfy-format save: `save_model` auto-keeps quantized storage
+      (comfy_quant markers) for convrot8 / nvfp4 / convrotcomfyw4a4 layers via
+      `toolkit/util/comfy_quant_export.py` (inverse of comfy_quant_import;
+      nvfp4 nibbles re-swapped + scales re-swizzled to the cuBLAS tile
+      layout), plain layers save at bf16; partially-exportable models fall
+      back to dequantized. Round-trip verified: save → mixin reload → outputs
+      match for convrot8, nvfp4, and plain layers.
+- [x] Save unification started: z_image and qwen_image holders now save
+      comfy-format single files via the mixin regardless of how they loaded
+      (z_image's dual-style branch deleted). Real round trip verified: the
+      published z_image int8_convrot file loads (270 quantized linears,
+      split-attach), resaves to the IDENTICAL 857-key comfy layout with
+      bit-exact fused qkv weights/scales/markers, and the reload's quantized
+      forward is bit-identical — toolkit saves are byte-compatible with
+      ComfyUI.
+- [x] chroma + chroma_radiance saves flipped to the mixin (their class keys
+      ARE the original layout) — also fixes their quanto-only dequant bug
+      (torchao/Ostris weights now dequantize on save). Tiny-model round trip
+      verified. Save flips so far: z_image, qwen_image, wan21, wan22_5b,
+      wan22_14b (dual files), chroma ×2.
+- [ ] flux_kontext comfy wiring deferred: its Comfy-Org repo ships a single
+      legacy-fp8 file in fused BFL layout — needs the flux fused-split
+      conversion (split_fused_quantized_keys pattern + BFL↔diffusers maps)
+- [ ] Flip the remaining per-arch `save_model` overrides as each arch's
+      save-side key conversion is in place
+- [ ] Publish/verify comfy repacks per model as they flip
+
+### Phase 3 — live server
+- [ ] Component-level identity (which TE/VAE instances are shared between archs) so
+      a model switch drops only what the next run doesn't need
+- [ ] Resident process: request comes in → diff requested components vs loaded →
+      unload/load the difference
+- [ ] Legacy `stable_diffusion_model.py` archs: grandfather or port last
+
+## 100% mixin coverage (2026-08-28)
+
+Every component of every non-legacy arch now loads through OstrisModelMixin —
+DiTs, text encoders, VAEs, vision encoders, connectors/vocoders, and the
+custom cases that previously bypassed it:
+
+- New wrapper classes: llama, gemma3 + gemma4, mistral3 (×2), qwen3 base,
+  qwen3-vl base + text-only, wan VAE, diffusers flux2 KL VAE, CLIP vision
+  (first vision encoder), cosmos DiT, anima text conditioner, the full ltx2
+  set (transformer, video/audio VAEs, connectors, vocoders).
+- Existing-class swaps: omnigen2 (mllm + VAE), boogu (TE + VAE), klein TE,
+  hidream llama, ideogram4 TE, prx TE.
+- Custom restructures: minimax video/audio VAEs and TE, MageVAE
+  (deferred-load ctor), the shared flux2_kl AutoEncoder (small-decoder sniff
+  as a class hook; flux2 + ideogram4 route through it), ace_step's bundle
+  (per-component class loaders), anima (modular-pipeline load replaced with
+  component-wise v2 loads + update_components), hidream_o1's Qwen3VL DiT,
+  z_image_l2p rebased onto the v2 class, ltx2's converter builds v2 classes.
+- Mixin: kwargs flow through the single-file chain (component ctor args like
+  MageVAE's sample_posterior).
+
+Verified with real weights: ltx2.3 + ltx2.5 (full v2 family), anima,
+ace_step_15 (new harness entry), minimax VAEs, plus the standing harness
+coverage. The legacy monolith archs joined the system per the inference-engine
+goal ("send a job with any base_model, reload/unload components on the fly"):
+six new wrapper classes complete the component vocabulary (UNet2DCondition,
+SD3/PixArt ×2/AuraFlow/Lumina2 transformers, Gemma2), and
+`adopt_component` (in-place class swap, OstrisLinear-style) rebinds
+pipeline-loaded components onto their wrappers at the monolith's single
+post-load funnel — covering every legacy arch without touching its fifteen
+load branches, with all pipeline references staying valid. Verified: sd1
+(UNet/KLVAE/CLIPTextEncoder all mixin instances + generation) and sdxl
+(dual-CLIP adoption + 1024² generation); harness gained sd1/sdxl entries,
+a legacy scheduler fallback, and the sampler-name pass-through. Full
+monolith decomposition (per-arch v2 loading with comfy candidates) remains
+future work, but every resident component is now poolable by the engine.
+
+## Unified load API (2026-08-28)
+
+Loading policy moved out of the holders into the mixin. The surface:
+
+- `ModelClassName.load(name_or_path, qtype=..., offload=..., dtype=...,
+  device=..., ...)` — sourcing (`load_model`) + `aitk_post_load` in one call.
+- `module.aitk_post_load(**kwargs)` — the post-load half alone, for holders
+  whose checkpoint sourcing is custom (state-dict surgery, combined files).
+  Handles: qtype (incl. `"qtype|ara_path"` accuracy recovery adapters, with
+  prequantized-skip via `aitk_is_quantized`), block-streamed quantization
+  (`quantize_module` in toolkit/util/quantize.py), MemoryManager layer
+  offloading with per-class `get_offload_ignore_modules()` hooks, and device
+  placement (low_vram parks on cpu).
+- `BaseModel.component_load_kwargs(role)` derives the kwargs from
+  model_config for roles "transformer" / "te" / "vae" (qtype/qtype_te, ARA
+  recombination, offload percents, low_vram, te/vae devices).
+
+All holders converted (transformer + TE quantize/offload/placement blocks
+deleted): krea2, z_image, qwen_image, chroma ×2, flux_kontext, f_light,
+nucleus, zeta_chroma, z_image_l2p, wan21 (+ subclasses), wan22_14b (dual:
+per-transformer load, dual-ARA branch kept), boogu, ernie, ideogram4,
+mageflow, prx_pixel, omnigen2, flux2 + klein, hidream + hidream_o1, ltx2 +
+ltx2.5 (prequantized ConvRot flagged via aitk_is_quantized), minimax_h3,
+anima (conditioner rides the transformer quantize flag at qtype_te),
+ace_step_15, example_model (template now teaches the new API).
+New `get_offload_ignore_modules` hooks: wan (scale_shift tables — also fixes
+wan21 offload which previously offloaded them), ltx2 (all four per-block
+tables), gemma3 (embed_tokens), gemma4 (embed_tokens + layer_scalar),
+ideogram4 (rotary inv_freq + input/cond projections), ernie (x_embedder),
+zeta ZImageDCT (pad tokens). New block-name hooks where only the holder had
+them: Flux2, MageFlow, ZImageDCT, Ideogram4Transformer2DModel (+ its exclude
+list on MageFlow). `prepare_text_encoder` remains only as the legacy-monolith
+path.
+
+## Load/quantize optimization round 1 (2026-08-28)
+
+Baseline: testing/.model_test_outputs/report_baseline.md (+ per-arch
+metrics_baseline.json). Profiling findings and fixes:
+
+- The "slow quantization" in the baseline was mostly NOT quantize math
+  (convrot8 on chroma's 57 blocks: 0.5s of GPU kernels). It was data
+  movement: mmap-backed safetensors faulting pages in per-tensor cudaMemcpy
+  order (cold cache on the 92%-full NVMe reads at ~0.32GB/s; the drive's
+  sequential ceiling is ~435MB/s — cold loads are storage-bound), plus the
+  old stream-quantize choreography crossing the bus three times (up, back
+  into a fresh host copy, up again at final placement).
+- `quantize_module(keep_on_device=True)`: when the final home IS the
+  quantize gpu (no offload/low_vram — aitk_post_load detects this), blocks
+  stay on the gpu after quantizing and the extras pass runs there too.
+  Weights cross the bus once; the model-sized host-RAM spike of the
+  intermediate copy is gone (chroma RSS peak 25.3 -> 17.3GB, the rest is
+  reclaimable page cache; flux2's 93GB spike class eliminated). Warm-cache
+  chroma blocks: 46s -> 3.9s; wan22_5b quantize 10.6s -> 3.8s.
+- Offload/low_vram path keeps cpu residency but quantizes extras layer-by-
+  layer via quantize_device instead of the all-cores cpu burst (was 350-650%
+  cpu avg with 1600% spikes).
+- `quantize()` skips the quantize_device round-trip for same-qtype
+  OstrisLinears (guaranteed no-op) and, for ostris backends, for non-Linear
+  leaves (norms/embeddings were being ferried across the bus for nothing).
+- posix_fadvise(WILLNEED) readahead on single-file loads (mixin._readahead)
+  — ~10-15% on cold reads, free when warm, no-op off POSIX.
+
+Round 2 (same day):
+- ARA path (`attach_ara_and_quantize`) now quantizes each hijacked linear on
+  the gpu (was: wherever it lived, historically the cpu) with the same
+  keep_on_device policy; the uint8 extras pass follows suit. Verified with
+  hidream + its uint3 ARA end to end.
+- `quantize_model` delegates its non-ARA branch to `quantize_module`
+  (keep_on_device derived from model_config); dead `_quantize_model_blocks`
+  removed; `patch_dequantization_on_save` made idempotent (double-patching
+  would have recursed on save).
+- Local-first hub loads: `_local_first` wraps from_pretrained / load_config /
+  AutoTokenizer / AutoProcessor — fully-cached repos load with
+  local_files_only (13x on VAE metadata, 0.53s -> 0.04s; a failed local
+  attempt is instant, so uncached repos fall through to the online path at
+  zero cost). Tokenizer-load slowness was DIAGNOSED as these per-process hub
+  etag HEADs, not the use_fast=False sites.
+
+Round 3 — vram spike elimination (2026-08-28):
+- keep_on_device's extras pass had moved the whole unquantized remainder to
+  the gpu in bf16 before quantizing (+7-10GB transient over final residency
+  on big-extras models). Now `quantize(keep_on_quantize_device=True)`
+  quantizes each extra layer on the gpu and leaves it — transient is one
+  layer; non-quantized leftovers move only at final placement. Same for the
+  ARA uint8 pass.
+- convrot8 `quantize_` chunked: takes the weight at stored dtype
+  (wants_fp32_weight=False; ConvRotIntN keeps True) and rotates+quantizes in
+  256MB row-chunks — a 152k-vocab TE projection's transient fell 7.6GB ->
+  1.4GB. Sub-threshold layers keep the one-shot path, verified bit-identical.
+- qwen_image measured: standard load peak 35.3 -> 27.5GB (below the 28.0
+  baseline; process peak now set by generation); low_vram load quantizes the
+  whole 20B stack never holding >2.4GB gpu.
+
+Ideas not yet done: pinned-buffer staged h2d (~+60% warm-upload bandwidth,
+9.5 vs 5.8GB/s measured), overlapping h2d with quantize kernels (small — the
+kernels are ~2% of the pass), wan22 low_vram generate choreography (165s of
+unload/reload per sample), holders' legacy "just to make sure" .to(device)
+lines that move quantized TEs to gpu even under low_vram (qwen's Preparing
+Model stage parks 7.6GB — policy question, not a quantize spike).
+
+## Offload hardening + final certification (2026-08-28)
+
+The harness gained a per-arch 100%-offload probe (attach MemoryManager at
+offload_percent=1.0 to the DiT + TEs + connectors/vision towers, then a
+2-step sample — 2 not 1: diffusers' shift_terminal stretch divides by zero
+at a single step). It flushed out and fixed, in the manager:
+- fully-offloaded modules now REPORT the compute device (class swap with
+  stringification-stable __module__/__qualname__ — transformers keys its
+  hidden-states capture registry by str(class)); bouncing ops also pull cpu
+  activations to the staged device
+- weights tied to embeddings (lm_head <-> embed_tokens) are never pinned out
+  from under the unmanaged embedding
+- ignore_modules are moved to the compute device at attach
+- the nested attach walk no longer double-lists managed modules as unmanaged
+  (a managed Embedding was getting hauled back to gpu by pipeline.to)
+- NEW: bouncing embeddings — vocab tables >64MB stay cpu-resident with the
+  row gather done ON cpu (only looked-up rows cross the bus); tied lm_heads
+  fall through to normal bouncing
+- wan22's dual wrapper skips its whole-14B .to() swaps when memory-managed,
+  and offload attaches per sub-transformer (pipelines hold those directly)
+
+Measured offload-phase peaks: wan22_14b 14.7 -> 1.5GB, ltx2.5 14.0 ->
+11.8GB (remainder: video activations + resident VAEs/vocoder by design),
+most image archs 1-4GB vs 10-53GB resident. flux2's Mistral warnings
+silenced behavior-preservingly (tie_word_embeddings=False, explicit
+fix_mistral_regex=False, processor_kwargs — tokenization verified
+bit-identical). ConvRot skip notices gated to layers >=1M weights.
+
+FINAL CERTIFICATION: full sweep, 33/34 PASS, 0 FAIL, 1 SKIP (mageflow
+upstream 404) — load + convrot8 quantize + generate + 100%-offload forward
+for every arch on the complete optimized stack.
+
+## Testing
+
+- [x] `testing/test_model_loading.py`: per-arch load + one small sample through
+      the normal training-style flow (get_model_class → load_model →
+      generate_images). `--arch X` runs one in-process; `--all` runs every
+      registered arch in its own subprocess (full unload between archs).
+      15 archs registered so far — add each model type as it migrates.
+- [x] Missing weights skip rather than fail: default is HF_HUB_OFFLINE=1 and
+      hub/file errors classify as SKIP; `--allow-download` opts into fetching.
+      (GPU + local-weights test, not CI-portable.)
+- [x] Final certification sweep (2026-08-27, post-polish): 14/14 runnable
+      archs PASS — comfy-source loads (zimage convrot8, qwen fp8, wan ×2 +
+      fp8 umt5 TE), all ported holder-config DiTs, and the migrated
+      quantize_model paths (chroma, flux_kontext, f_light block-streamed) in
+      one run; mageflow remains the upstream 404 skip. One regression caught
+      and fixed: qwen's _load_single_file override needed the new config
+      kwarg.
+- [x] Full sweep run 2026-08-27: 14/15 PASS (zimage, qwen_image, krea2,
+      boogu_image, ernie_image, ideogram4, hidream_o1, anima, wan21, wan22_5b,
+      chroma, flux_kontext, flux2_klein_4b, ltx2.3 — the quantized 22B ltx
+      stack doesn't fit 32GB, needs the 96GB card). mageflow blocked
+      upstream: microsoft/Mage-Flow-Base 404s on the hub (cached locally, so
+      it runs offline — recheck whether the repo moved/went private).
+- [x] Registry carries realistic per-arch sample settings (native res, steps,
+      CFG) so sweep outputs are visually verifiable, not just "a file
+      exists". Verified: all 14 produce proper generations. Findings from
+      the quality pass: boogu emits a black frame below native res at
+      low-step/high-CFG (settings regime, present pre-restructure, not a
+      migration bug); chroma's FakeCLIP hardcoded device 'cuda' broke any
+      non-cuda:0 run (pre-existing, fixed — FakeCLIP now takes the real
+      device); ideogram4's fp8 release renders its own "blocked by safety
+      filter" card for a plain cat prompt (model behavior, not a bug —
+      investigate its trigger).
+- [x] Round-trip verified for the first comfy-save arch: z_image convrot
+      load → comfy save → identical key set + bit-exact quantized entries vs
+      the published file → reload → bit-identical quantized forward. Extend
+      per arch as saves flip.
+- [x] Full UI-default coverage (2026-08-27): the registry now holds all 31
+      UI-facing archs, and 30/31 load + generate through the new stack with
+      their real UI defaults (mageflow blocked upstream by its hub 404).
+      This round verified the previously-untested tail with real weights:
+      wan22_14b + i2v (comfy fp8 pairs via the legacy importer, candidate
+      keys added for the ai-toolkit bf16 default repos), wan21_i2v, hidream,
+      hidream_e1 (native 768² editing), nucleus, omnigen2, ltx2.5, flux2,
+      flux2_klein_9b, prx_pixel, zeta_chroma, zimage_l2p, both qwen edit
+      archs, f-lite. Fixes found by the run: ltx2.5's fp32 scale_shift
+      tables promoted hidden states into bf16 linears under the diffusers
+      class — fixed ComfyUI-style (toolkit/util/mixed_precision.py):
+      per-op input casting hooks on every weighted module + the stored-fp32
+      tensors pinned against parent .to(dtype) casts (device moves pass
+      through), so the tables stay genuinely fp32 at sample time. The
+      mechanism is general — any mixed-precision comfy checkpoint on a
+      diffusers-class arch can use it; harness configs
+      for e1 resolution and zeta/l2p extras_name_or_path corrected to match
+      the UI defaults.
+- [ ] Each newly migrated model adds its test in the same PR as its migration.
+
+## TODO / look at later
+
+- [x] Quantize consolidation: `quantize_model` now honors `quantize_kwargs`
+      (blocks + extras) and tolerates missing block names; the chroma ×2,
+      flux_kontext, f_light, omnigen2 raw-quantize sites migrated onto it
+      (gaining block streaming, excludes, ARA, dequant-on-save patching) with
+      holder block names added. Remaining raw sites are legacy/extension
+      (flex2, cogview4, stable_diffusion_model). The ARA uint8 hardcode
+      stands — revisit if a non-uint8 ARA base is ever wanted.
+- [x] Last known `qtype_te` bugs fixed (flux2's Mistral TE, anima's
+      text_conditioner) — 9/9 sites from the survey now correct outside the
+      grandfathered legacy monolith (cogview4/legacy SD remain as-is).
+- [x] wan comfy-TE resolved for real: UMT5TextEncoder carries comfy
+      candidates (fp8_e4m3fn_scaled via the legacy importer, fp16), files
+      already in transformers key layout (spiece blob dropped, tied
+      embed_tokens materialized). Verified: wan21 samples with the local
+      comfy fp8 TE. The loaders/umt5.py `comfy_files` param stays as a
+      no-op shim for old callers.
+- [x] Registry hardening: unknown archs now raise with the known-arch list
+      (legacy monolith archs whitelisted via LEGACY_ARCHS). Still open: lazy
+      per-arch imports; single source of truth shared with the UI's
+      `options.tsx` model list.
+- [x] Stub dedup where identical: chroma_radiance imports FakeCLIP/FakeConfig
+      from chroma_model. The other Fake* copies (hidream_o1, flux2, zeta)
+      carry model-specific values — left in place.
+- [ ] Vendored upstream code (hidream/src, omnigen2/src, ltx2 converter's private
+      comfy-quant parser): dedupe against toolkit utils where practical.

toolkit/models/v2/_mixin.py

@@ -9,7 +9,11 @@ digests any of:
   - a HuggingFace repo id ("org/repo")
   - a local single .safetensors file in the model's original key layout
   - a remote single .safetensors file ("org/repo/path/file.safetensors")
-plus optional automatic quantization of the loaded weights (`qtype=...`).
+plus optional automatic quantization of the loaded weights (`qtype=...`). A
+single-file checkpoint carrying ComfyUI `comfy_quant` markers loads with its
+pre-quantized layers attached to the toolkit's quantization backends
+(convrot8 / nvfp4) automatically. `save_model` writes a single-file
+.safetensors back in the original (comfy) key layout.
 
 Model specific behavior lives in small overridable hooks (key conversion, config
 source, block names, backend loading), so subclassing for a new model usually means
@@ -38,6 +42,24 @@ class BasicModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMix
     pass
 
 
+def adopt_component(module: torch.nn.Module, wrapper_cls: type) -> torch.nn.Module:
+    """Rebind an already-loaded component instance onto its v2 wrapper class
+    (in-place class swap, like the OstrisLinear conversion). Used when a
+    component arrives from machinery that builds the base class directly
+    (e.g. a diffusers pipeline load), so every resident component is a mixin
+    instance the inference engine can pool and manage."""
+    if isinstance(module, wrapper_cls):
+        return module
+    base = wrapper_cls.__mro__[1]
+    if not isinstance(module, base):
+        raise TypeError(
+            f"cannot adopt {type(module).__name__} into {wrapper_cls.__name__} "
+            f"(expected a {base.__name__})"
+        )
+    module.__class__ = wrapper_cls
+    return module
+
+
 class OstrisModelMixin:
     # ---- per-model configuration, override in subclasses ----
     # subfolder that holds this model inside a diffusers style checkpoint
@@ -47,6 +69,35 @@ class OstrisModelMixin:
     # .safetensors file and no config_path is given
     aitk_config_repo: Optional[str] = None
 
+    # ---- comfy weight sources ----
+    # hub repo the comfy-format weight files are published in (Comfy-Org/...)
+    aitk_comfy_repo: Optional[str] = None
+    # standard name_or_path (hub repo id) -> repo-relative comfy weight file
+    # candidates for this module: every precision variant the class can digest.
+    # Selection ranks them convrot8 > float8 mixed > float8 > bf16 > fp16
+    # (resolver.comfy_precision_rank); a locally-present candidate always wins
+    # over a download.
+    aitk_comfy_weight_names: Dict[str, List[str]] = {}
+
+    # ---- tokenizer/processor source, for text-encoder modules ----
+    aitk_tokenizer_repo: Optional[str] = None
+    aitk_tokenizer_subfolder: Optional[str] = None
+    aitk_processor_repo: Optional[str] = None
+    aitk_processor_subfolder: Optional[str] = None
+
+    # single-file loads without quant markers cast tensors to the requested
+    # dtype; classes whose checkpoints carry a deliberate precision mix (e.g.
+    # fp32 norms next to bf16 weights) set this False to load at stored
+    # precision instead
+    aitk_cast_on_load: bool = True
+
+    # pre-quantized (comfy marker) checkpoints normally load their
+    # NON-quantized tensors at stored precision (the mix can be deliberate,
+    # e.g. ltx2.5's fp32 tables). Classes whose forward assumes one uniform
+    # dtype (wan: fp16 conv biases next to fp32 tables in the fp8 files) set
+    # this True to cast those float tensors to the load dtype instead
+    aitk_cast_quantized_load: bool = False
+
     # ---- state set by the loader / quantizer ----
     aitk_is_quantized: bool = False
     aitk_qtype: Optional[str] = None
@@ -68,7 +119,7 @@ class OstrisModelMixin:
         return state_dict
 
     @classmethod
-    def get_quantization_block_names(cls) -> Optional[List[str]]:
+    def get_transformer_block_names(cls) -> Optional[List[str]]:
         """Names (dotted paths allowed) of the repeated block lists to quantize one
         block at a time so the whole model never has to sit on the gpu at once."""
         return None
@@ -78,28 +129,201 @@ class OstrisModelMixin:
         """fnmatch patterns of sensitive modules to keep in full precision."""
         return None
 
+    def get_offload_ignore_modules(self) -> Optional[List]:
+        """Module/parameter objects layer offloading must keep resident (pad
+        tokens, modulation tables, ...). Default: none."""
+        return None
+
     # ---- backend hooks: default to the diffusers ModelMixin API. transformers-lib
     # models (text encoders) override these three to use PreTrainedModel/AutoConfig.
 
     @classmethod
     def aitk_from_pretrained(cls, path, subfolder=None, dtype=None, **kwargs):
-        return cls.from_pretrained(
-            path, subfolder=subfolder, torch_dtype=dtype, **kwargs
+        return cls._local_first(
+            cls.from_pretrained, path, subfolder=subfolder, torch_dtype=dtype, **kwargs
         )
 
     @classmethod
     def aitk_load_config(cls, path, subfolder=None):
-        return cls.load_config(path, subfolder=subfolder)
+        return cls._local_first(cls.load_config, path, subfolder=subfolder)
+
+    @staticmethod
+    def _local_first(loader, *args, **kwargs):
+        """Serve fully-cached repos without touching the hub. Online
+        from_pretrained does an etag HEAD round-trip per file (seconds of pure
+        network latency on every warm load); a local_files_only attempt is
+        instant when the cache can't satisfy it, so failing over to the
+        online path costs nothing."""
+        import os as _os
+
+        if _os.environ.get("HF_HUB_OFFLINE") == "1":
+            return loader(*args, **kwargs)
+        try:
+            return loader(*args, local_files_only=True, **kwargs)
+        except Exception:
+            return loader(*args, **kwargs)
 
     @classmethod
     def aitk_from_config(cls, config):
-        with torch.device("meta"):
-            return cls.from_config(config)
+        # params on meta (materialized by the state-dict assign), buffers real:
+        # non-persistent buffers (e.g. rope tables) are computed at init and
+        # never appear in checkpoints
+        from accelerate import init_empty_weights
+
+        with init_empty_weights(include_buffers=False):
+            if hasattr(cls, "from_config"):
+                return cls.from_config(config)
+            # plain nn.Module classes take their params/config object directly
+            return cls(config)
 
     # ------------------------------------------------------------------
     # loading
     # ------------------------------------------------------------------
 
+    @classmethod
+    def load(
+        cls,
+        name_or_path: str,
+        qtype: Optional[str] = None,
+        offload: float = 0.0,
+        dtype: torch.dtype = torch.bfloat16,
+        device: Optional[torch.device] = None,
+        base_model=None,
+        quantize_device: Optional[torch.device] = None,
+        exclude_quant_modules: Optional[List[str]] = None,
+        config_path: Optional[str] = None,
+        config=None,
+        subfolder: Optional[str] = None,
+        use_comfy_weights: bool = True,
+        **kwargs,
+    ):
+        """One-call component load. Resolves/downloads the source (comfy
+        candidates included), builds the module, quantizes it, attaches layer
+        offloading, and places it — everything a holder used to hand-roll.
+
+        qtype: quantization type; "type|path.safetensors" also attaches an
+        accuracy recovery adapter (requires base_model for the network
+        wiring). Pre-quantized checkpoints skip quantization automatically.
+        offload: layer-offload fraction (0 disables); the class's
+        get_offload_ignore_modules hook pins sensitive pieces.
+        device: final placement. With offload active the memory manager owns
+        staging, so no blanket device move happens.
+        base_model: the holder — supplies status prints, quantize_kwargs, and
+        ARA wiring when given.
+        """
+        model = cls.load_model(
+            name_or_path,
+            dtype=dtype,
+            config_path=config_path,
+            config=config,
+            subfolder=subfolder,
+            use_comfy_weights=use_comfy_weights,
+            **kwargs,
+        )
+        return model.aitk_post_load(
+            qtype=qtype,
+            offload=offload,
+            dtype=dtype,
+            device=device,
+            base_model=base_model,
+            quantize_device=quantize_device,
+            exclude_quant_modules=exclude_quant_modules,
+        )
+
+    def aitk_post_load(
+        self,
+        qtype: Optional[str] = None,
+        offload: float = 0.0,
+        dtype: torch.dtype = torch.bfloat16,
+        device: Optional[torch.device] = None,
+        base_model=None,
+        quantize_device: Optional[torch.device] = None,
+        exclude_quant_modules: Optional[List[str]] = None,
+        **_load_only_kwargs,  # tolerate component_load_kwargs() extras
+    ):
+        """The post-load half of `load` — quantize (incl. "qtype|ara"), attach
+        layer offloading, place on device. Callable directly by holders whose
+        checkpoint sourcing is custom but want the standard policy."""
+        from toolkit.print import print_acc
+
+        status_fn = (
+            base_model.print_and_status_update if base_model is not None else print_acc
+        )
+
+        ara_path = None
+        if qtype is not None and "|" in qtype:
+            qtype, ara_path = qtype.split("|", 1)
+
+        if qtype and not getattr(self, "aitk_is_quantized", False):
+            from toolkit.util.quantize import (
+                attach_ara_and_quantize,
+                quantize_module,
+            )
+
+            exclude = list(exclude_quant_modules or []) + list(
+                self.get_quantization_exclude_modules() or []
+            )
+            q_device = quantize_device if quantize_device is not None else device
+            # final home is the quantize gpu (no offloading, no cpu
+            # parking): quantized weights stay put — one bus crossing
+            # instead of three, and no model-sized host-ram copy
+            keep_on_device = (
+                not offload
+                and device is not None
+                and q_device is not None
+                and torch.device(device) == torch.device(q_device)
+                and torch.device(device).type != "cpu"
+            )
+            if ara_path is not None:
+                if base_model is None:
+                    raise ValueError(
+                        "an accuracy recovery adapter needs base_model= for the "
+                        "network wiring"
+                    )
+                status_fn("Quantizing with accuracy recovery adapter")
+                attach_ara_and_quantize(
+                    base_model,
+                    self,
+                    ara_path,
+                    exclude=exclude,
+                    device=q_device,
+                    keep_on_device=keep_on_device,
+                )
+            else:
+                status_fn(f"Quantizing ({qtype})")
+                quantize_module(
+                    self,
+                    qtype,
+                    device=q_device,
+                    dtype=dtype,
+                    block_names=self.get_transformer_block_names(),
+                    exclude=exclude,
+                    quantize_kwargs=(
+                        base_model.model_config.quantize_kwargs
+                        if base_model is not None
+                        else None
+                    ),
+                    status_fn=status_fn,
+                    keep_on_device=keep_on_device,
+                )
+            self.aitk_is_quantized = True
+            self.aitk_qtype = qtype
+        elif qtype and getattr(self, "aitk_is_quantized", False):
+            status_fn("Checkpoint is pre-quantized; skipping quantization")
+
+        if offload and offload > 0:
+            from toolkit.memory_management import MemoryManager
+
+            MemoryManager.attach(
+                self,
+                device,
+                offload_percent=offload,
+                ignore_modules=list(self.get_offload_ignore_modules() or []),
+            )
+        elif device is not None:
+            self.to(device)
+        return self
+
     @classmethod
     def load_model(
         cls,
@@ -110,7 +334,9 @@ class OstrisModelMixin:
         quantize_device: Optional[torch.device] = None,
         exclude_quant_modules: Optional[List[str]] = None,
         config_path: Optional[str] = None,
+        config=None,
         subfolder: Optional[str] = None,
+        use_comfy_weights: bool = True,
         **kwargs,
     ):
         """Load a model universally from a given name or path.
@@ -119,19 +345,47 @@ class OstrisModelMixin:
         single .safetensors file, or a remote single file
         ("org/repo/file.safetensors").
 
+        When name_or_path is a standard hub repo this class has comfy weight
+        candidates registered for (aitk_comfy_weight_names), the comfy file is
+        the preferred source: the best-ranked local candidate, or the
+        top-preference candidate downloaded into the comfy layout under
+        MODELS_PATH. The standard repo still supplies the config. Local dirs
+        and unregistered repos load as-is; use_comfy_weights=False opts out.
+
         qtype: quantize the weights after loading. quantize_device: where to run the
         quantization math; blocks are moved there one at a time and returned to where
         they were.
         config_path: config source for single-file loads, overriding aitk_config_repo.
         device: move the finished model there before returning.
+        subfolder: overrides the class's aitk_subfolder; pass "" to explicitly
+        load from the checkpoint root (e.g. a raw hub repo).
         """
         if subfolder is None:
             subfolder = cls.aitk_subfolder
+        elif subfolder == "":
+            subfolder = None
+
+        if (
+            use_comfy_weights
+            and not name_or_path.endswith(".safetensors")
+            and not os.path.exists(name_or_path)
+        ):
+            comfy_path = cls.resolve_comfy_weights(name_or_path, subfolder=subfolder)
+            if comfy_path is not None:
+                if config_path is None:
+                    # the standard repo supplies the config for the comfy file
+                    config_path = name_or_path
+                name_or_path = comfy_path
 
         if name_or_path.endswith(".safetensors"):
             file_path = cls._resolve_single_file(name_or_path)
             model = cls._load_single_file(
-                file_path, dtype=dtype, config_path=config_path, subfolder=subfolder
+                file_path,
+                dtype=dtype,
+                config_path=config_path,
+                config=config,
+                subfolder=subfolder,
+                **kwargs,
             )
         else:
             if os.path.isdir(name_or_path):
@@ -188,28 +442,296 @@ class OstrisModelMixin:
             subfolder = None
         return cls.aitk_load_config(config_source, subfolder=subfolder)
 
+    @staticmethod
+    def _readahead(file_path: str):
+        """Queue kernel readahead for the whole file. safetensors tensors are
+        mmap-backed; without this a cold cache faults pages in per-tensor
+        cudaMemcpy order at a fraction of the drive's sequential bandwidth."""
+        if not hasattr(os, "posix_fadvise"):
+            return  # non-POSIX: mmap readahead heuristics only
+        try:
+            fd = os.open(file_path, os.O_RDONLY)
+            try:
+                os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_WILLNEED)
+            finally:
+                os.close(fd)
+        except OSError:
+            pass
+
     @classmethod
     def _load_single_file(
         cls,
         file_path: str,
         dtype: torch.dtype,
         config_path: Optional[str] = None,
+        config=None,
         subfolder: Optional[str] = None,
+        **kwargs,
     ):
-        config = cls._load_single_file_config(config_path, subfolder)
-
+        cls._readahead(file_path)
         state_dict = load_file(file_path)
-        state_dict = cls.convert_state_dict_on_load(state_dict)
-        for key, value in state_dict.items():
-            state_dict[key] = value.to(dtype=dtype)
+        return cls.load_from_state_dict(
+            state_dict,
+            dtype,
+            config_path=config_path,
+            config=config,
+            subfolder=subfolder,
+            **kwargs,
+        )
 
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict: Dict[str, torch.Tensor]):
+        """Derive the model config from the checkpoint itself (e.g. sniffing
+        block counts from key indices). Return None (the default) to load the
+        config from config_path / aitk_config_repo instead."""
+        return None
+
+    @classmethod
+    def load_from_state_dict(
+        cls,
+        state_dict: Dict[str, torch.Tensor],
+        dtype: torch.dtype,
+        config_path: Optional[str] = None,
+        config=None,
+        subfolder: Optional[str] = None,
+        **kwargs,
+    ):
+        """Build the model and load an already-read single-file state dict
+        (the tail of the single-file path; also callable directly for
+        checkpoints read from non-safetensors sources). ``config``, when
+        given, is used directly (for models whose config comes from the
+        holder, e.g. model_kwargs-driven archs)."""
+        state_dict = cls.convert_state_dict_on_load(state_dict)
+        has_quant_markers = "scaled_fp8" in state_dict or any(
+            k.endswith(".comfy_quant") for k in state_dict
+        )
+        if config is None:
+            config = cls.aitk_config_from_state_dict(state_dict)
+        if config is None:
+            config = cls._load_single_file_config(config_path, subfolder)
         model = cls.aitk_from_config(config)
-        model.load_state_dict(state_dict, assign=True)
-        model.to(dtype=dtype)
+
+        if has_quant_markers:
+            # pre-quantized comfy checkpoint: attach the quantized linears onto
+            # the toolkit's backends, load the rest at its stored precision
+            # (the checkpoint's bf16/fp16/fp32 mix is deliberate)
+            from toolkit.util.comfy_quant_import import import_comfy_quantized_layers
+
+            state_dict, num_quantized = import_comfy_quantized_layers(
+                model, state_dict, orig_dtype=dtype
+            )
+            if cls.aitk_cast_quantized_load:
+                # uniform-dtype models: the leftover (non-quantized) float
+                # tensors follow the load dtype, like the non-marker path
+                for key, value in state_dict.items():
+                    if value.is_floating_point():
+                        state_dict[key] = value.to(dtype=dtype)
+            cls._load_state_dict_with_quantized(model, state_dict)
+            if cls.aitk_cast_quantized_load:
+                # the importer assigns quantized-layer biases directly at
+                # stored dtype; they must follow too (a stray fp16 bias also
+                # makes ModelMixin.dtype — the pipelines' cast target — lie)
+                from toolkit.util.ostris_quant import OstrisLinear
+
+                for m in model.modules():
+                    if isinstance(m, OstrisLinear) and m.bias is not None:
+                        m.bias.data = m.bias.data.to(dtype=dtype)
+            model.aitk_is_quantized = True
+        elif cls.aitk_cast_on_load:
+            for key, value in state_dict.items():
+                if value.is_floating_point():
+                    state_dict[key] = value.to(dtype=dtype)
+            model.load_state_dict(state_dict, assign=True)
+            model.to(dtype=dtype)
+        else:
+            # stored-precision load (the checkpoint's dtype mix is deliberate)
+            model.load_state_dict(state_dict, assign=True)
         del state_dict
         flush()
         return model
 
+    @staticmethod
+    def _load_state_dict_with_quantized(model, state_dict):
+        """Load a state dict onto a (meta-built) model whose quantized linears
+        were already attached by import_comfy_quantized_layers: their weight
+        (and importer-assigned bias) legitimately report as missing keys."""
+        from toolkit.util.ostris_quant import OstrisLinear
+
+        result = model.load_state_dict(state_dict, assign=True, strict=False)
+        quantized_param_keys = set()
+        for name, m in model.named_modules():
+            if isinstance(m, OstrisLinear):
+                quantized_param_keys.add(f"{name}.weight")
+                if m.bias is not None:
+                    quantized_param_keys.add(f"{name}.bias")
+        bad_missing = [k for k in result.missing_keys if k not in quantized_param_keys]
+        if bad_missing or result.unexpected_keys:
+            raise ValueError(
+                f"{type(model).__name__} load mismatch: missing {bad_missing[:8]}, "
+                f"unexpected {result.unexpected_keys[:8]}"
+            )
+        leftover_meta = [n for n, p in model.named_parameters() if p.is_meta]
+        if leftover_meta:
+            raise ValueError(
+                f"{type(model).__name__} load left meta parameters: "
+                f"{leftover_meta[:8]}"
+            )
+
+    # ------------------------------------------------------------------
+    # comfy weight sources
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def resolve_comfy_weights(
+        cls,
+        name_or_path: str,
+        subfolder: Optional[str] = None,
+        local_only: bool = False,
+        hf_token: Optional[str] = None,
+        status_fn: Optional[callable] = None,
+    ) -> Optional[str]:
+        """The comfy-format weight file replacing a standard ``name_or_path``,
+        or None when this class has none registered for it. Best-ranked local
+        candidate wins; otherwise the top-preference candidate is downloaded
+        to the comfy layout under MODELS_PATH (unless local_only).
+
+        Candidate keys may be plain repo ids or ``(repo_id, subfolder)``
+        tuples for checkpoints holding several of this component (e.g.
+        wan2.2's transformer / transformer_2)."""
+        from toolkit.models.v2.resolver import resolve_comfy_candidates
+
+        candidates = None
+        if subfolder is not None:
+            candidates = cls.aitk_comfy_weight_names.get((name_or_path, subfolder))
+        if candidates is None:
+            candidates = cls.aitk_comfy_weight_names.get(name_or_path)
+        repo_id = cls.aitk_comfy_repo
+        if isinstance(candidates, dict):
+            # entries may carry their own comfy repo ({"repo": ..., "files": [...]})
+            repo_id = candidates.get("repo", repo_id)
+            candidates = candidates.get("files")
+        if not candidates or repo_id is None:
+            return None
+        return resolve_comfy_candidates(
+            candidates,
+            repo_id=repo_id,
+            hf_token=hf_token,
+            status_fn=status_fn,
+            local_only=local_only,
+        )
+
+    # ------------------------------------------------------------------
+    # tokenizer / processor (text-encoder modules)
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def load_tokenizer(
+        cls,
+        name_or_path: Optional[str] = None,
+        subfolder: Optional[str] = None,
+        **kwargs,
+    ):
+        """Load the tokenizer from ``name_or_path`` (a checkpoint dir or repo
+        holding it at aitk_tokenizer_subfolder), falling back to the class's
+        aitk_tokenizer_repo. subfolder overrides the class default; "" loads
+        from the checkpoint root."""
+        from transformers import AutoTokenizer
+
+        source = name_or_path if name_or_path is not None else cls.aitk_tokenizer_repo
+        if source is None:
+            raise ValueError(f"{cls.__name__} does not declare aitk_tokenizer_repo")
+        if subfolder is None:
+            subfolder = cls.aitk_tokenizer_subfolder
+        if subfolder and os.path.isdir(source) and not os.path.isdir(
+            os.path.join(source, subfolder)
+        ):
+            subfolder = None
+        return cls._local_first(
+            AutoTokenizer.from_pretrained, source, subfolder=subfolder or "", **kwargs
+        )
+
+    @classmethod
+    def load_processor(
+        cls,
+        name_or_path: Optional[str] = None,
+        subfolder: Optional[str] = None,
+        **kwargs,
+    ):
+        from transformers import AutoProcessor
+
+        source = name_or_path if name_or_path is not None else cls.aitk_processor_repo
+        if source is None:
+            raise ValueError(f"{cls.__name__} does not declare aitk_processor_repo")
+        if subfolder is None:
+            subfolder = cls.aitk_processor_subfolder
+        if subfolder and os.path.isdir(source) and not os.path.isdir(
+            os.path.join(source, subfolder)
+        ):
+            subfolder = None
+        return cls._local_first(
+            AutoProcessor.from_pretrained, source, subfolder=subfolder or "", **kwargs
+        )
+
+    # ------------------------------------------------------------------
+    # saving
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def save_model(
+        self,
+        output_path: str,
+        dtype: Optional[torch.dtype] = None,
+        metadata: Optional[Dict[str, str]] = None,
+        quantized: Optional[bool] = None,
+    ):
+        """Save as a single-file .safetensors in the model's original (comfy)
+        key layout, via convert_state_dict_on_save.
+
+        quantized: None (auto) keeps quantized storage — comfy_quant markers,
+        loadable by ComfyUI and by this class — when the model holds quantized
+        layers whose backend has a comfy format (convrot8 / nvfp4 /
+        convrotcomfyw4a4), and saves dequantized full precision otherwise.
+        True forces quantized storage (raises if any quantized layer's backend
+        has no comfy format); False forces a dequantized save. dtype, when
+        given, casts the floating point (non-quantized) tensors."""
+        from safetensors.torch import save_file
+        from toolkit.util.quantize import dequantize_if_quantized
+
+        q_entries: Dict[str, torch.Tensor] = {}
+        exported: List[str] = []
+        if quantized is None or quantized:
+            from toolkit.util.comfy_quant_export import export_comfy_quantized_layers
+
+            q_entries, exported, unexportable = export_comfy_quantized_layers(self)
+            if unexportable:
+                if quantized:
+                    raise ValueError(
+                        f"{type(self).__name__}: quantized layers without a comfy "
+                        f"storage format: {unexportable[:8]}"
+                    )
+                # auto mode: a partially-quantized file would be inconsistent —
+                # fall back to a fully dequantized save
+                q_entries, exported = {}, []
+
+        skip_keys = {f"{name}.weight" for name in exported}
+        state_dict = {}
+        for key, value in self.state_dict().items():
+            if key in skip_keys:
+                continue
+            value = dequantize_if_quantized(value)
+            if dtype is not None and value.is_floating_point():
+                value = value.to(dtype=dtype)
+            state_dict[key] = value.detach().to("cpu").contiguous()
+        state_dict.update(q_entries)
+        state_dict = self.convert_state_dict_on_save(state_dict)
+
+        parent = os.path.dirname(output_path)
+        if parent:
+            os.makedirs(parent, exist_ok=True)
+        save_file(state_dict, output_path, metadata=metadata)
+        del state_dict
+        flush()
+
     # ------------------------------------------------------------------
     # quantization
     # ------------------------------------------------------------------
@@ -222,7 +744,7 @@ class OstrisModelMixin:
         exclude: Optional[List[str]] = None,
     ):
         """Quantize the model weights in place. When device is given, the repeated
-        blocks (get_quantization_block_names) are moved there one at a time for the
+        blocks (get_transformer_block_names) are moved there one at a time for the
         quantization math and returned to their original device, so the whole model
         never has to fit on the gpu in full precision."""
         from optimum.quanto import freeze
@@ -238,7 +760,7 @@ class OstrisModelMixin:
         )
 
         blocks: List[torch.nn.Module] = []
-        for name in self.get_quantization_block_names() or []:
+        for name in self.get_transformer_block_names() or []:
             # name may be a dotted path for models that nest their blocks
             block_list = self
             for part in name.split("."):
@@ -270,3 +792,30 @@ class OstrisModelMixin:
         self.aitk_qtype = qtype
         flush()
         return self
+
+
+class OstrisTransformersMixin(OstrisModelMixin):
+    """OstrisModelMixin with the backend hooks speaking the transformers-lib
+    API (PreTrainedModel / AutoConfig) instead of diffusers ModelMixin. Base
+    for text-encoder and vision-encoder modules."""
+
+    @classmethod
+    def aitk_from_pretrained(cls, path, subfolder=None, dtype=None, **kwargs):
+        return cls._local_first(
+            cls.from_pretrained, path, subfolder=subfolder or "", torch_dtype=dtype, **kwargs
+        )
+
+    @classmethod
+    def aitk_load_config(cls, path, subfolder=None):
+        from transformers import AutoConfig
+
+        return cls._local_first(
+            AutoConfig.from_pretrained, path, subfolder=subfolder or ""
+        )
+
+    @classmethod
+    def aitk_from_config(cls, config):
+        from accelerate import init_empty_weights
+
+        with init_empty_weights(include_buffers=False):
+            return cls(config)

toolkit/models/v2/diffusion_models/auraflow.py

@@ -0,0 +1,13 @@
+from diffusers import AuraFlowTransformer2DModel as DiffusersAuraFlowTransformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class AuraFlowTransformer2DModel(
+    DiffusersAuraFlowTransformer2DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["joint_transformer_blocks", "single_transformer_blocks"]

toolkit/models/v2/diffusion_models/cosmos.py

@@ -0,0 +1,15 @@
+from diffusers.models import (
+    CosmosTransformer3DModel as DiffusersCosmosTransformer3DModel,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class CosmosTransformer3DModel(DiffusersCosmosTransformer3DModel, OstrisModelMixin):
+    """Cosmos video/image DiT (anima)."""
+
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]

toolkit/models/v2/diffusion_models/flux.py

@@ -0,0 +1,14 @@
+from diffusers import FluxTransformer2DModel as DiffusersFluxTransformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class FluxTransformer2DModel(DiffusersFluxTransformer2DModel, OstrisModelMixin):
+    """Flux1-family DiT (flux, flux_kontext, chroma-adjacent finetunes in
+    diffusers layout)."""
+
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks", "single_transformer_blocks"]

toolkit/models/v2/diffusion_models/hidream.py

@@ -0,0 +1,18 @@
+from diffusers.models import (
+    HiDreamImageTransformer2DModel as DiffusersHiDreamImageTransformer2DModel,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class HiDreamImageTransformer2DModel(
+    DiffusersHiDreamImageTransformer2DModel, OstrisModelMixin
+):
+    """The diffusers HiDream DiT (hidream_e1; the base hidream arch uses the
+    vendored copy in the hidream extension)."""
+
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["double_stream_blocks", "single_stream_blocks"]

toolkit/models/v2/diffusion_models/ltx2.py

@@ -0,0 +1,46 @@
+from diffusers.models.transformers import (
+    LTX2VideoTransformer3DModel as DiffusersLTX2VideoTransformer3DModel,
+)
+from diffusers.pipelines.ltx2 import (
+    LTX2TextConnectors as DiffusersLTX2TextConnectors,
+)
+from diffusers.pipelines.ltx2 import LTX2Vocoder as DiffusersLTX2Vocoder
+from diffusers.pipelines.ltx2 import (
+    LTX2VocoderWithBWE as DiffusersLTX2VocoderWithBWE,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class LTX2VideoTransformer3DModel(
+    DiffusersLTX2VideoTransformer3DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]
+
+    def get_offload_ignore_modules(self):
+        # fp32 scale/shift tables (bare tensors) must stay resident
+        mods = []
+        for block in self.transformer_blocks:
+            mods += [
+                block.scale_shift_table,
+                block.audio_scale_shift_table,
+                block.video_a2v_cross_attn_scale_shift_table,
+                block.audio_a2v_cross_attn_scale_shift_table,
+            ]
+        return mods + [self.scale_shift_table, self.audio_scale_shift_table]
+
+
+class LTX2TextConnectors(DiffusersLTX2TextConnectors, OstrisModelMixin):
+    aitk_subfolder = "connectors"
+
+
+class LTX2Vocoder(DiffusersLTX2Vocoder, OstrisModelMixin):
+    aitk_subfolder = "vocoder"
+
+
+class LTX2VocoderWithBWE(DiffusersLTX2VocoderWithBWE, OstrisModelMixin):
+    aitk_subfolder = "vocoder"

toolkit/models/v2/diffusion_models/lumina2.py

@@ -0,0 +1,13 @@
+from diffusers import Lumina2Transformer2DModel as DiffusersLumina2Transformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class Lumina2Transformer2DModel(
+    DiffusersLumina2Transformer2DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]

toolkit/models/v2/diffusion_models/nucleus_image.py

@@ -0,0 +1,15 @@
+from diffusers import (
+    NucleusMoEImageTransformer2DModel as DiffusersNucleusMoEImageTransformer2DModel,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class NucleusMoEImageTransformer2DModel(
+    DiffusersNucleusMoEImageTransformer2DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]

toolkit/models/v2/diffusion_models/pixart.py

@@ -0,0 +1,24 @@
+from diffusers import PixArtTransformer2DModel as DiffusersPixArtTransformer2DModel
+from diffusers import Transformer2DModel as DiffusersTransformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class PixArtTransformer2DModel(
+    DiffusersPixArtTransformer2DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]
+
+
+class Transformer2DModel(DiffusersTransformer2DModel, OstrisModelMixin):
+    """The generic diffusers DiT the pixart sigma path loads."""
+
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]

toolkit/models/v2/diffusion_models/qwen_image.py

@@ -0,0 +1,58 @@
+from diffusers import (
+    QwenImageTransformer2DModel as DiffusersQwenImageTransformer2DModel,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class QwenImageTransformer2DModel(
+    DiffusersQwenImageTransformer2DModel, OstrisModelMixin
+):
+    aitk_subfolder = "transformer"
+    aitk_config_repo = "Qwen/Qwen-Image"
+
+    aitk_comfy_repo = "Comfy-Org/Qwen-Image_ComfyUI"
+    # the comfy files use the diffusers key layout directly (no conversion);
+    # fp8mixed carries float8_e4m3fn comfy_quant markers that attach straight
+    # onto this class's modules
+    aitk_comfy_weight_names = {
+        "Qwen/Qwen-Image": [
+            "split_files/diffusion_models/qwen_image_fp8mixed.safetensors",
+            # raw fp8 cast (no markers, diffusers keys) — loads via from_single_file
+            "split_files/diffusion_models/qwen_image_fp8_e4m3fn.safetensors",
+            "split_files/diffusion_models/qwen_image_bf16.safetensors",
+        ],
+    }
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]
+
+    @classmethod
+    def _load_single_file(
+        cls, file_path, dtype, config_path=None, config=None, subfolder=None
+    ):
+        from safetensors import safe_open
+
+        with safe_open(file_path, framework="pt") as f:
+            has_markers = any(k.endswith(".comfy_quant") for k in f.keys())
+        if has_markers:
+            # comfy prequantized checkpoint (diffusers key layout): the mixin
+            # path attaches the quantized layers
+            return super()._load_single_file(
+                file_path,
+                dtype,
+                config_path=config_path,
+                config=config,
+                subfolder=subfolder,
+            )
+        # other single-file checkpoints carry diffusers or original key
+        # layouts; diffusers' single-file machinery owns that conversion
+        model = cls.from_single_file(
+            file_path,
+            config=config_path if config_path is not None else cls.aitk_config_repo,
+            subfolder="transformer",
+            torch_dtype=dtype,
+        )
+        model.to(dtype)
+        return model

toolkit/models/v2/diffusion_models/sd3.py

@@ -0,0 +1,11 @@
+from diffusers import SD3Transformer2DModel as DiffusersSD3Transformer2DModel
+
+from .._mixin import OstrisModelMixin
+
+
+class SD3Transformer2DModel(DiffusersSD3Transformer2DModel, OstrisModelMixin):
+    aitk_subfolder = "transformer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]

toolkit/models/v2/diffusion_models/unet.py

@@ -0,0 +1,13 @@
+from diffusers import UNet2DConditionModel as DiffusersUNet2DConditionModel
+
+from .._mixin import OstrisModelMixin
+
+
+class UNet2DConditionModel(DiffusersUNet2DConditionModel, OstrisModelMixin):
+    """The SD1/SD2/SDXL-family UNet."""
+
+    aitk_subfolder = "unet"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["down_blocks", "up_blocks"]

toolkit/models/v2/diffusion_models/wan.py

@@ -0,0 +1,204 @@
+import torch
+import torch.nn.functional as F
+from diffusers import WanTransformer3DModel as DiffusersWanTransformer3DModel
+
+from .._mixin import OstrisModelMixin
+
+
+def patch_wan_patch_embedding(transformer) -> bool:
+    """Wan's patch_embedding is a Conv3d whose kernel == stride (non-overlapping
+    patchify), 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 the forward for the exactly-equivalent F.linear (a GEMM),
+    returning the conv's own [b, out, t', h', w'] layout so it stays a drop-in.
+    The weight is read lazily so this survives later .to(device)/dtype moves."""
+    conv = getattr(transformer, "patch_embedding", None)
+    if not isinstance(conv, torch.nn.Conv3d):
+        return False
+    if (
+        tuple(conv.kernel_size) != tuple(conv.stride)
+        or any(p != 0 for p in conv.padding)
+        or any(d != 1 for d in conv.dilation)
+        or conv.groups != 1
+    ):
+        return False
+
+    def fast_forward(x, _conv=conv):
+        weight = _conv.weight
+        out_c = weight.shape[0]
+        kt, kh, kw = _conv.kernel_size
+        b, c, t, h, w = x.shape
+        tp, hp, wp = t // kt, h // kh, w // kw
+        # patchify with (c, kt, kh, kw) fastest — matches weight.reshape below
+        x = x.reshape(b, c, tp, kt, hp, kh, wp, kw)
+        x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(b, tp * hp * wp, -1)
+        w2d = weight.reshape(out_c, -1)
+        x = F.linear(x.to(w2d.dtype), w2d, _conv.bias)
+        return x.transpose(1, 2).reshape(b, out_c, tp, hp, wp)
+
+    conv.forward = fast_forward
+    return True
+
+
+class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin):
+    """Wan 2.1/2.2 video DiT (wan22 loads two of these into its dual wrapper)."""
+
+    aitk_subfolder = "transformer"
+
+    aitk_comfy_repo = "Comfy-Org/Wan_2.2_ComfyUI_Repackaged"
+    # comfy wan files use the original key layout; convert_state_dict_on_load
+    # renames them (pure substring renames, so the legacy scaled_fp8 weight /
+    # scale keys ride along with their modules). wan2.2 A14B checkpoints hold
+    # two DiTs, keyed by (repo, subfolder): transformer = high noise,
+    # transformer_2 = low noise. wan2.1 entries override the comfy repo.
+    aitk_comfy_weight_names = {
+        "Wan-AI/Wan2.2-TI2V-5B-Diffusers": [
+            "split_files/diffusion_models/wan2.2_ti2v_5B_fp16.safetensors",
+        ],
+        ("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "transformer"): [
+            "split_files/diffusion_models/wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors",
+        ],
+        ("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "transformer_2"): [
+            "split_files/diffusion_models/wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors",
+        ],
+        ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer"): [
+            "split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors",
+        ],
+        ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer_2"): [
+            "split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors",
+        ],
+        # the UI defaults point at the ai-toolkit bf16 repacks; same comfy files
+        ("ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "transformer"): [
+            "split_files/diffusion_models/wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors",
+        ],
+        ("ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "transformer_2"): [
+            "split_files/diffusion_models/wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors",
+        ],
+        ("ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "transformer"): [
+            "split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors",
+        ],
+        ("ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "transformer_2"): [
+            "split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors",
+        ],
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": {
+            "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+            "files": [
+                "split_files/diffusion_models/wan2.1_t2v_1.3B_bf16.safetensors",
+                "split_files/diffusion_models/wan2.1_t2v_1.3B_fp16.safetensors",
+            ],
+        },
+        "Wan-AI/Wan2.1-T2V-14B-Diffusers": {
+            "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+            "files": [
+                "split_files/diffusion_models/wan2.1_t2v_14B_fp8_scaled.safetensors",
+                "split_files/diffusion_models/wan2.1_t2v_14B_bf16.safetensors",
+                "split_files/diffusion_models/wan2.1_t2v_14B_fp16.safetensors",
+            ],
+        },
+        "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": {
+            "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+            "files": [
+                "split_files/diffusion_models/wan2.1_i2v_480p_14B_fp8_scaled.safetensors",
+                "split_files/diffusion_models/wan2.1_i2v_480p_14B_bf16.safetensors",
+                "split_files/diffusion_models/wan2.1_i2v_480p_14B_fp16.safetensors",
+            ],
+        },
+        "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": {
+            "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+            "files": [
+                "split_files/diffusion_models/wan2.1_i2v_720p_14B_fp8_scaled.safetensors",
+                "split_files/diffusion_models/wan2.1_i2v_720p_14B_bf16.safetensors",
+                "split_files/diffusion_models/wan2.1_i2v_720p_14B_fp16.safetensors",
+            ],
+        },
+    }
+
+    # the comfy fp8/scaled_fp8 wan files mix fp16 conv biases with fp32
+    # tables; wan's forward assumes one uniform dtype, so cast the
+    # non-quantized tensors to the load dtype (the old holder did this with a
+    # blanket .to(device, dtype) after load)
+    aitk_cast_quantized_load = True
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["blocks"]
+
+    def aitk_post_load(self, **kwargs):
+        # slow bf16 Conv3d patchify -> equivalent GEMM, before quantize/offload
+        patch_wan_patch_embedding(self)
+        return super().aitk_post_load(**kwargs)
+
+    def get_offload_ignore_modules(self):
+        # fp32 modulation tables must stay resident on the compute device
+        return [self.scale_shift_table] + [
+            block.scale_shift_table for block in self.blocks
+        ]
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        # original/comfy wan keys -> diffusers layout via diffusers' own
+        # converter (rename-only for the base/i2v/vace variants)
+        is_original = any(
+            ".self_attn." in k
+            or ".cross_attn." in k
+            or k.startswith("model.diffusion_model.")
+            for k in state_dict
+        )
+        if not is_original:
+            return state_dict
+        from diffusers.loaders.single_file_utils import (
+            convert_wan_transformer_to_diffusers,
+        )
+
+        return convert_wan_transformer_to_diffusers(dict(state_dict))
+
+    # inverse of diffusers' rename table, for the base/t2v/i2v variants (no
+    # vace/animate: their extra mappings collide in reverse). Order matters:
+    # longest/most-specific first, and the norm2/norm3 swap uses a placeholder.
+    _SAVE_RENAMES = [
+        ("condition_embedder.time_embedder.linear_1", "time_embedding.0"),
+        ("condition_embedder.time_embedder.linear_2", "time_embedding.2"),
+        ("condition_embedder.text_embedder.linear_1", "text_embedding.0"),
+        ("condition_embedder.text_embedder.linear_2", "text_embedding.2"),
+        ("condition_embedder.time_proj", "time_projection.1"),
+        ("condition_embedder.image_embedder.norm1", "img_emb.proj.0"),
+        ("condition_embedder.image_embedder.ff.net.0.proj", "img_emb.proj.1"),
+        ("condition_embedder.image_embedder.ff.net.2", "img_emb.proj.3"),
+        ("condition_embedder.image_embedder.norm2", "img_emb.proj.4"),
+        ("ffn.net.0.proj", "ffn.0"),
+        ("ffn.net.2", "ffn.2"),
+        (".norm_added_k.", ".norm_k_img."),
+        (".add_k_proj.", ".k_img."),
+        (".add_v_proj.", ".v_img."),
+        (".to_out.0.", ".o."),
+        (".to_q.", ".q."),
+        (".to_k.", ".k."),
+        (".to_v.", ".v."),
+        ("attn2", "cross_attn"),
+        ("attn1", "self_attn"),
+        # norm2 <-> norm3 swap back
+        ("norm3", "norm__placeholder"),
+        ("norm2", "norm3"),
+        ("norm__placeholder", "norm2"),
+        ("proj_out", "head.head"),
+    ]
+
+    @classmethod
+    def convert_state_dict_on_save(cls, state_dict):
+        """Diffusers layout back to the original/comfy wan key layout
+        (rename-only, so quantized weight/scale/marker keys ride along)."""
+        if any(".self_attn." in k or ".cross_attn." in k for k in state_dict):
+            return state_dict  # already original layout
+        new_sd = {}
+        for key, value in state_dict.items():
+            k = key
+            # scale_shift_table: blocks keep the name as `modulation`, the
+            # top-level one belongs to the output head
+            if k == "scale_shift_table":
+                k = "head.modulation"
+            elif k.endswith(".scale_shift_table"):
+                k = k[: -len("scale_shift_table")] + "modulation"
+            for src, dst in cls._SAVE_RENAMES:
+                k = k.replace(src, dst)
+            new_sd[k] = value
+        return new_sd

toolkit/models/v2/diffusion_models/z_image.py

@@ -3,7 +3,7 @@ from diffusers.models.transformers import (
     ZImageTransformer2DModel as DiffusersZImageTransformer2DModel,
 )
 
-from ._mixin import OstrisModelMixin
+from .._mixin import OstrisModelMixin
 
 
 class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMixin):
@@ -11,10 +11,24 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
     # repo to pull the config from when loading a single-file checkpoint
     aitk_config_repo = "Tongyi-MAI/Z-Image-Turbo"
 
+    aitk_comfy_repo = "Comfy-Org/z_image_turbo"
+    # the int8_convrot file marks the FUSED attention.qkv modules; the load
+    # converter splits those entries exactly into to_q/to_k/to_v (row-sliced
+    # qdata + scales), so it attaches onto this split layout directly
+    aitk_comfy_weight_names = {
+        "Tongyi-MAI/Z-Image-Turbo": [
+            "split_files/diffusion_models/z_image_turbo_int8_convrot.safetensors",
+            "split_files/diffusion_models/z_image_turbo_bf16.safetensors",
+        ],
+    }
+
     @classmethod
-    def get_quantization_block_names(cls):
+    def get_transformer_block_names(cls):
         return ["layers"]
 
+    def get_offload_ignore_modules(self):
+        return [self.x_pad_token, self.cap_pad_token]
+
     @classmethod
     def get_quantization_exclude_modules(cls):
         # sensitive modules kept in full precision (fnmatch patterns on module
@@ -36,6 +50,22 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
     @classmethod
     def convert_state_dict_on_load(cls, state_dict):
         """Convert a single-file Z-Image checkpoint to diffusers transformer keys."""
+        from toolkit.util.comfy_quant_import import split_fused_quantized_keys
+
+        state_dict = dict(state_dict)
+        # quantized fused qkv entries (comfy convrot/fp8/nvfp4 files) split
+        # exactly into the three projections: row-consecutive weights/scales
+        for marker_key in [
+            k for k in list(state_dict) if k.endswith(".attention.qkv.comfy_quant")
+        ]:
+            prefix = marker_key[: -len(".comfy_quant")]
+            base = prefix[: -len(".qkv")]
+            split_fused_quantized_keys(
+                state_dict,
+                prefix,
+                [f"{base}.to_q", f"{base}.to_k", f"{base}.to_v"],
+            )
+
         new_sd = {}
         for key, value in state_dict.items():
             k = key
@@ -47,9 +77,10 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
                 new_sd[prefix + ".attention.to_k.weight"] = k_proj
                 new_sd[prefix + ".attention.to_v.weight"] = v
                 continue
-            k = k.replace(".attention.out.weight", ".attention.to_out.0.weight")
-            k = k.replace(".attention.q_norm.weight", ".attention.norm_q.weight")
-            k = k.replace(".attention.k_norm.weight", ".attention.norm_k.weight")
+            # module-prefix renames so weight/bias/scale/marker keys all map
+            k = k.replace(".attention.out.", ".attention.to_out.0.")
+            k = k.replace(".attention.q_norm.", ".attention.norm_q.")
+            k = k.replace(".attention.k_norm.", ".attention.norm_k.")
             if k.startswith("x_embedder."):
                 k = "all_x_embedder.2-1." + k[len("x_embedder.") :]
             elif k.startswith("final_layer."):
@@ -60,6 +91,20 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
     @classmethod
     def convert_state_dict_on_save(cls, state_dict):
         """Convert a diffusers transformer state dict back to the single-file layout."""
+        from toolkit.util.comfy_quant_import import fuse_split_quantized_keys
+
+        state_dict = dict(state_dict)
+        # quantized split projections fuse back into the single-file qkv entry
+        for marker_key in [
+            k for k in list(state_dict) if k.endswith(".attention.to_q.comfy_quant")
+        ]:
+            base = marker_key[: -len(".to_q.comfy_quant")]
+            fuse_split_quantized_keys(
+                state_dict,
+                [f"{base}.to_q", f"{base}.to_k", f"{base}.to_v"],
+                f"{base}.qkv",
+            )
+
         new_sd = {}
         qkv_cache = {}
         for key, value in state_dict.items():
@@ -90,9 +135,9 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
                     break
             if matched:
                 continue
-            k = k.replace(".attention.to_out.0.weight", ".attention.out.weight")
-            k = k.replace(".attention.norm_q.weight", ".attention.q_norm.weight")
-            k = k.replace(".attention.norm_k.weight", ".attention.k_norm.weight")
+            k = k.replace(".attention.to_out.0.", ".attention.out.")
+            k = k.replace(".attention.norm_q.", ".attention.q_norm.")
+            k = k.replace(".attention.norm_k.", ".attention.k_norm.")
             if k.startswith("all_x_embedder.2-1."):
                 k = "x_embedder." + k[len("all_x_embedder.2-1.") :]
             elif k.startswith("all_final_layer.2-1."):

toolkit/models/v2/resolver.py

@@ -0,0 +1,201 @@
+"""ComfyUI-layout weight file resolution.
+
+Weight files live under MODELS_PATH in ComfyUI's folder layout
+(diffusion_models/, text_encoders/, vae/, ...) so the folder is shareable with
+a ComfyUI install. Files are used in place when present and downloaded to
+exactly their repo-relative location only when missing, so nothing is ever
+duplicated on re-run.
+
+Lifted from the minimax_h3 / ltx2.5 model implementations; those now call
+into here.
+"""
+
+import os
+from typing import Callable, Iterable, Optional
+
+from toolkit.paths import MODELS_PATH
+
+
+def comfy_precision_rank(filename: str) -> int:
+    """Load-preference rank for a comfy weight filename:
+    convrot8 (0) > float8 mixed (1) > float8 (2) > bf16 (3) > fp16 (4) >
+    anything else, e.g. nvfp4 or unmarked (5)."""
+    name = os.path.basename(filename).lower()
+    if "convrot" in name:
+        return 0
+    is_fp8 = "fp8" in name or "float8" in name or "e4m3" in name
+    if is_fp8 and "mixed" in name:
+        return 1
+    if is_fp8:
+        return 2
+    if "bf16" in name:
+        return 3
+    if "fp16" in name:
+        return 4
+    return 5
+
+
+def comfy_local_rel(repo_rel: str) -> str:
+    """Repo file path -> ComfyUI models-folder path. Comfy-Org repos nest the
+    comfy layout under a packaging prefix (split_files/, non_official/) that is
+    not part of the shared models folder layout."""
+    for prefix in ("split_files/", "non_official/"):
+        if repo_rel.startswith(prefix):
+            return repo_rel[len(prefix):]
+    return repo_rel
+
+
+def resolve_comfy_candidates(
+    candidates: Iterable[str],
+    repo_id: str,
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+    local_only: bool = False,
+) -> Optional[str]:
+    """Pick the best comfy weight file among precision variants of one
+    component (repo-relative paths, ranked by comfy_precision_rank then list
+    order). The best-ranked LOCAL candidate wins; only when no candidate is
+    local is the best-ranked one downloaded to its comfy-layout location
+    under MODELS_PATH."""
+    candidates = list(candidates)
+    ordered = sorted(
+        candidates, key=lambda c: (comfy_precision_rank(c), candidates.index(c))
+    )
+    for repo_rel in ordered:
+        found = resolve_comfy_file(
+            comfy_local_rel(repo_rel), repo_id, local_only=True
+        )
+        if found is not None:
+            return found
+    if local_only:
+        return None
+
+    import huggingface_hub
+
+    best = ordered[0]
+    local_rel = comfy_local_rel(best)
+    if status_fn is not None:
+        status_fn(f"Downloading {best} from {repo_id} into {MODELS_PATH}")
+    path = huggingface_hub.hf_hub_download(
+        repo_id=repo_id, filename=best, token=hf_token, local_dir=MODELS_PATH
+    )
+    target = os.path.join(MODELS_PATH, local_rel)
+    if os.path.abspath(path) != os.path.abspath(target):
+        # move out of the packaging prefix into the shared comfy layout
+        os.makedirs(os.path.dirname(target), exist_ok=True)
+        os.replace(path, target)
+        return target
+    return path
+
+
+def find_file_recursive(root_dir: str, filename: str) -> Optional[str]:
+    """First (breadth-stable, sorted) match of ``filename`` anywhere under
+    ``root_dir``."""
+    if not os.path.isdir(root_dir):
+        return None
+    for dirpath, dirnames, filenames in os.walk(root_dir):
+        dirnames.sort()
+        if filename in filenames:
+            return os.path.join(dirpath, filename)
+    return None
+
+
+def repo_id_from_name_or_path(
+    name_or_path: Optional[str], default: str
+) -> str:
+    """Treat a hub-style ``name_or_path`` ("org/repo") as a replacement comfy
+    repo; anything local (or an explicit .safetensors file) keeps the
+    default."""
+    if (
+        name_or_path
+        and not os.path.exists(name_or_path)
+        and not name_or_path.endswith(".safetensors")
+        and "/" in name_or_path
+    ):
+        return name_or_path
+    return default
+
+
+def resolve_comfy_file(
+    rel_path: str,
+    repo_id: str,
+    override_path: Optional[str] = None,
+    extra_roots: Optional[Iterable[str]] = None,
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+    local_only: bool = False,
+) -> Optional[str]:
+    """Find a weight file at its local location, or download it there when
+    (and only when) it is missing.
+
+    Search order: ``override_path`` (must exist), the repo-relative path under
+    MODELS_PATH (and each of ``extra_roots``), the bare filename at each root,
+    any subfolder of the category folder (recursive — e.g.
+    diffusion_models/my_custom_sub/), then the hub — downloaded to the
+    repo-relative path under MODELS_PATH. With ``local_only`` the hub is never
+    touched and a miss returns None.
+    """
+    if override_path is not None:
+        if not os.path.exists(override_path):
+            raise FileNotFoundError(
+                f"Override path for {rel_path} does not exist: {override_path}"
+            )
+        return override_path
+
+    filename = os.path.basename(rel_path)
+    category = os.path.dirname(rel_path)
+    roots = [MODELS_PATH] + [r for r in (extra_roots or []) if os.path.isdir(r)]
+    for root in roots:
+        for rel in (rel_path, filename):
+            candidate = os.path.join(root, rel)
+            if os.path.exists(candidate):
+                return candidate
+    for root in roots:
+        found = find_file_recursive(os.path.join(root, category), filename)
+        if found is not None:
+            return found
+
+    if local_only:
+        return None
+
+    import huggingface_hub
+
+    if status_fn is not None:
+        status_fn(f"Downloading {rel_path} from {repo_id} into {MODELS_PATH}")
+    return huggingface_hub.hf_hub_download(
+        repo_id=repo_id, filename=rel_path, token=hf_token, local_dir=MODELS_PATH
+    )
+
+
+def resolve_named_file(
+    path: str,
+    component: str = "model",
+    hf_token: Optional[str] = None,
+) -> str:
+    """Resolve an explicit .safetensors reference: a local file, a file already
+    under MODELS_PATH, or an 'org/repo/path/file.safetensors' hub path
+    (downloaded into the models folder at its repo-relative path)."""
+    if os.path.exists(path):
+        return path
+    splits = path.split("/")
+    if len(splits) < 3:
+        raise ValueError(
+            f"Invalid {component} path: {path}. Must be a local file or "
+            "'org/repo/filename.safetensors' to download from the Hugging Face Hub."
+        )
+    rel_path = "/".join(splits[2:])
+    for candidate in (
+        os.path.join(MODELS_PATH, rel_path),
+        os.path.join(MODELS_PATH, splits[-1]),
+    ):
+        if os.path.exists(candidate):
+            return candidate
+
+    import huggingface_hub
+
+    return huggingface_hub.hf_hub_download(
+        repo_id="/".join(splits[:2]),
+        filename=rel_path,
+        token=hf_token,
+        local_dir=MODELS_PATH,
+    )

toolkit/models/v2/text_encoders/anima.py

@@ -0,0 +1,9 @@
+from diffusers import AnimaTextConditioner as DiffusersAnimaTextConditioner
+
+from .._mixin import OstrisModelMixin
+
+
+class AnimaTextConditioner(DiffusersAnimaTextConditioner, OstrisModelMixin):
+    """Anima's learned text conditioner (rides next to the Qwen3 encoder)."""
+
+    aitk_subfolder = "text_conditioner"

toolkit/models/v2/text_encoders/clip.py

@@ -0,0 +1,27 @@
+from transformers import CLIPTextModel, CLIPTextModelWithProjection
+
+from .._mixin import OstrisTransformersMixin
+
+
+class CLIPTextEncoder(CLIPTextModel, OstrisTransformersMixin):
+    """CLIP-L text encoder (flux-style checkpoints: text_encoder/ +
+    tokenizer/)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["text_model.encoder.layers"]
+
+
+class CLIPTextEncoderWithProjection(CLIPTextModelWithProjection, OstrisTransformersMixin):
+    """CLIP text encoder with the projection head (SDXL / SD3 / HiDream style
+    checkpoints; the second encoder lives at text_encoder_2/ + tokenizer_2/)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["text_model.encoder.layers"]

toolkit/models/v2/text_encoders/gemma2.py

@@ -0,0 +1,14 @@
+from transformers import Gemma2Model
+
+from .._mixin import OstrisTransformersMixin
+
+
+class Gemma2ModelEncoder(Gemma2Model, OstrisTransformersMixin):
+    """Gemma2 base model (lumina2's text encoder — what AutoModel resolves)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]

toolkit/models/v2/text_encoders/gemma3.py

@@ -0,0 +1,42 @@
+from transformers import Gemma3ForConditionalGeneration
+
+from .._mixin import OstrisTransformersMixin
+
+
+class Gemma3TextEncoder(Gemma3ForConditionalGeneration, OstrisTransformersMixin):
+    """Gemma3 conditioning stack (ltx2 / ltx2.3)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        # both layouts seen across transformers versions; missing paths skip
+        return ["model.language_model.layers", "language_model.model.layers"]
+
+    def get_offload_ignore_modules(self):
+        return [self.model.language_model.base_model.embed_tokens]
+
+
+try:
+    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel
+
+    class Gemma4TextEncoder(Gemma4TextModel, OstrisTransformersMixin):
+        """Gemma4 text decoder (ltx2.5's conditioning stack)."""
+
+        aitk_subfolder = "text_encoder"
+        aitk_tokenizer_subfolder = "tokenizer"
+
+        @classmethod
+        def get_transformer_block_names(cls):
+            return ["layers"]
+
+        def get_offload_ignore_modules(self):
+            # layer_scalar is a bare tensor buffer on each decoder layer; the
+            # manager never enumerates it, so it must ride along explicitly
+            return [self.embed_tokens] + [
+                layer.layer_scalar for layer in self.layers
+            ]
+
+except ImportError:
+    Gemma4TextEncoder = None

toolkit/models/v2/text_encoders/llama.py

@@ -0,0 +1,14 @@
+from transformers import LlamaForCausalLM
+
+from .._mixin import OstrisTransformersMixin
+
+
+class LlamaTextEncoder(LlamaForCausalLM, OstrisTransformersMixin):
+    """Llama causal-LM text encoder (hidream's text_encoder_4)."""
+
+    aitk_subfolder = "text_encoder_4"
+    aitk_tokenizer_subfolder = "tokenizer_4"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["model.layers"]

toolkit/models/v2/text_encoders/mistral3.py

@@ -0,0 +1,25 @@
+from transformers import Mistral3ForConditionalGeneration, Mistral3Model
+
+from .._mixin import OstrisTransformersMixin
+
+
+class Mistral3TextEncoder(Mistral3ForConditionalGeneration, OstrisTransformersMixin):
+    """Mistral3 conditioning stack (flux2)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["model.language_model.layers", "language_model.model.layers"]
+
+
+class Mistral3ModelEncoder(Mistral3Model, OstrisTransformersMixin):
+    """The inner Mistral3 base model (ernie_image's text encoder)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["language_model.layers"]

toolkit/models/v2/text_encoders/qwen25_vl.py

@@ -0,0 +1,22 @@
+from transformers import Qwen2_5_VLForConditionalGeneration
+
+from .._mixin import OstrisTransformersMixin
+
+
+class Qwen25VLTextEncoder(Qwen2_5_VLForConditionalGeneration, OstrisTransformersMixin):
+    """Qwen2.5-VL conditioning stack (qwen_image family). Loads from a
+    checkpoint's text_encoder/ subfolder; the edit variants keep the vision
+    tower, plain t2i drops it."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["model.language_model.layers"]
+
+    def drop_vision_tower(self):
+        """Text-only conditioning: the vision tower is dead weight."""
+        if getattr(self.model, "visual", None) is not None:
+            self.model.visual = None
+        return self

toolkit/models/v2/text_encoders/qwen3.py

@@ -0,0 +1,28 @@
+from transformers import Qwen3ForCausalLM, Qwen3Model
+
+from .._mixin import OstrisTransformersMixin
+
+
+class Qwen3TextEncoder(Qwen3ForCausalLM, OstrisTransformersMixin):
+    """Qwen3 causal-LM text encoder (Z-Image family, Zeta-Chroma, ...). Loads
+    from a checkpoint's text_encoder/ subfolder, a hub repo, or a single
+    .safetensors file; the tokenizer rides in the checkpoint's tokenizer/
+    subfolder."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["model.layers"]
+
+
+class Qwen3ModelEncoder(Qwen3Model, OstrisTransformersMixin):
+    """The inner Qwen3 base model (anima's text encoder)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]

toolkit/models/v2/text_encoders/qwen3_vl.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn.functional as F
+from transformers import (
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLModel,
+    Qwen3VLTextModel,
+)
+
+from .._mixin import OstrisTransformersMixin
+
+
+def patch_qwen_vl_patch_embed(model) -> int:
+    """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."""
+    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
+
+
+class Qwen3VLTextEncoder(Qwen3VLForConditionalGeneration, OstrisTransformersMixin):
+    """Qwen3-VL conditioning stack (krea2, mageflow, nucleus_image, ideogram4,
+    minimax_h3, ...). Loads from a checkpoint's text_encoder/ subfolder or a
+    raw Qwen repo (pass subfolder=\"\" for the latter)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_processor_subfolder = "processor"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["model.language_model.layers"]
+
+    def drop_vision_tower(self):
+        """Text-only conditioning: the vision tower is dead weight — drop it to
+        free VRAM and skip loading its (bf16-slow) Conv3d patch_embed."""
+        if getattr(self.model, "visual", None) is not None:
+            self.model.visual = None
+        return self
+
+    def patch_vision_patch_embed(self) -> int:
+        """Keep the vision tower (reference images ride into the embeddings)
+        but swap its Conv3d patch_embed for an equivalent GEMM."""
+        return patch_qwen_vl_patch_embed(self)
+
+
+class Qwen3VLModelEncoder(Qwen3VLModel, OstrisTransformersMixin):
+    """The inner Qwen3-VL base model (what AutoModel resolves): its
+    last_hidden_state is the instruction feature boogu_image / ideogram4
+    consume."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["language_model.layers"]
+
+    def drop_vision_tower(self):
+        if getattr(self, "visual", None) is not None:
+            self.visual = None
+        return self
+
+    def patch_vision_patch_embed(self) -> int:
+        return patch_qwen_vl_patch_embed(self)
+
+
+class Qwen3VLTextOnlyEncoder(Qwen3VLTextModel, OstrisTransformersMixin):
+    """The Qwen3-VL text tower alone (prx_pixel)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]

toolkit/models/v2/text_encoders/t5.py

@@ -0,0 +1,16 @@
+from transformers import T5EncoderModel
+
+from .._mixin import OstrisTransformersMixin
+
+
+class T5TextEncoder(T5EncoderModel, OstrisTransformersMixin):
+    """T5-XXL text encoder. Defaults to the flux-style checkpoint layout
+    (text_encoder_2/ + tokenizer_2/); pass subfolder overrides for checkpoints
+    that keep it at text_encoder/ + tokenizer/ (e.g. f-lite)."""
+
+    aitk_subfolder = "text_encoder_2"
+    aitk_tokenizer_subfolder = "tokenizer_2"
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["encoder.block"]

toolkit/models/v2/text_encoders/umt5.py

@@ -0,0 +1,77 @@
+import torch
+from transformers import T5Tokenizer, UMT5EncoderModel
+
+from .._mixin import OstrisTransformersMixin
+
+
+class PatchedT5Tokenizer(T5Tokenizer):
+    def __init__(
+        self,
+        vocab=None,
+        eos_token="</s>",
+        unk_token="<unk>",
+        pad_token="<pad>",
+        _spm_precompiled_charsmap=None,
+        extra_ids=100,
+        additional_special_tokens=None,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab=vocab,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            _spm_precompiled_charsmap=None,  # this is passing a empty byte string for some reason now.
+            extra_ids=extra_ids,
+            additional_special_tokens=additional_special_tokens,
+            **kwargs,
+        )
+
+
+class UMT5TextEncoder(UMT5EncoderModel, OstrisTransformersMixin):
+    """UMT5-XXL text encoder (wan family)."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+    aitk_config_repo = "ai-toolkit/umt5_xxl_encoder"
+    # comfy fp8 umt5 files store non-quantized tensors in a dtype mix; T5
+    # assumes a uniform dtype, so follow the load dtype
+    aitk_cast_quantized_load = True
+
+    aitk_comfy_repo = "Comfy-Org/Wan_2.1_ComfyUI_repackaged"
+    # comfy umt5 files already use the transformers key layout; the
+    # fp8_e4m3fn_scaled variant is the legacy scaled-fp8 format (handled by
+    # the importer's float8 backend)
+    aitk_comfy_weight_names = {
+        "ai-toolkit/umt5_xxl_encoder": [
+            "split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors",
+            "split_files/text_encoders/umt5_xxl_fp16.safetensors",
+        ],
+    }
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["encoder.block"]
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        # drop the embedded sentencepiece blob and materialize the tied
+        # embed_tokens reference
+        state_dict = dict(state_dict)
+        state_dict.pop("spiece_model", None)
+        if (
+            "encoder.embed_tokens.weight" not in state_dict
+            and "shared.weight" in state_dict
+        ):
+            state_dict["encoder.embed_tokens.weight"] = state_dict["shared.weight"]
+        return state_dict
+
+    @classmethod
+    def load_tokenizer(cls, name_or_path=None, subfolder=None, **kwargs):
+        # T5's tokenizer needs the _spm_precompiled_charsmap patch
+        source = name_or_path if name_or_path is not None else cls.aitk_tokenizer_repo
+        if subfolder is None:
+            subfolder = cls.aitk_tokenizer_subfolder
+        return PatchedT5Tokenizer.from_pretrained(
+            source, subfolder=subfolder or "", **kwargs
+        )

toolkit/models/v2/vae/autoencoder_kl.py

@@ -0,0 +1,10 @@
+from diffusers import AutoencoderKL
+
+from .._mixin import OstrisModelMixin
+
+
+class KLVAE(AutoencoderKL, OstrisModelMixin):
+    """The diffusers AutoencoderKL (SD/SDXL/Flux1/Z-Image image VAEs), loaded
+    from a checkpoint's vae/ subfolder through the universal loader."""
+
+    aitk_subfolder = "vae"

toolkit/models/v2/vae/autoencoder_kl_flux2.py

@@ -0,0 +1,10 @@
+from diffusers import AutoencoderKLFlux2
+
+from .._mixin import OstrisModelMixin
+
+
+class Flux2KLVAE(AutoencoderKLFlux2, OstrisModelMixin):
+    """The diffusers Flux2 KL VAE (ernie_image; distinct from the hand-rolled
+    BFL AutoEncoder in vae/flux2_kl.py that flux2/ideogram4 use)."""
+
+    aitk_subfolder = "vae"

toolkit/models/v2/vae/flux2_kl.py

@@ -1,4 +1,13 @@
-"""Flux2 KL autoencoder."""
+"""The BFL-style Flux2 KL autoencoder (32ch latents, 2x2 pixel-shuffle
+packing to 128ch), shared by the flux2 family and ideogram4.
+
+flux2 loads the raw BFL ae.safetensors layout and uses encode/decode (with the
+BatchNorm running-stats latent normalization); ideogram4 loads diffusers-format
+checkpoints via convert_diffusers_state_dict and drives encoder/decoder
+directly with its own patchify + latent-norm tables. The two models pack the
+128 latent channels in different orders — the packing/normalization stays
+per-model.
+"""
 
 from __future__ import annotations
 
@@ -11,6 +20,8 @@ import torch.utils.checkpoint as ckpt
 from einops import rearrange
 from torch import Tensor, nn
 
+from .._mixin import OstrisModelMixin
+
 
 @dataclass
 class AutoEncoderParams:
@@ -22,6 +33,17 @@ class AutoEncoderParams:
     num_res_blocks: int = 2
     z_channels: int = 32
 
+@dataclass
+class AutoEncoderSmallDecoderParams:
+    resolution: int = 256
+    in_channels: int = 3
+    ch: int = 128
+    ch_encoder: int = 96
+    out_ch: int = 3
+    ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
+    num_res_blocks: int = 2
+    z_channels: int = 32
+
 
 def swish(x: Tensor) -> Tensor:
     return x * torch.sigmoid(x)
@@ -194,30 +216,30 @@ class Encoder(nn.Module):
         for i_level in range(self.num_resolutions):
             for i_block in range(self.num_res_blocks):
                 if torch.is_grad_enabled() and self.gradient_checkpointing:
-                    h = ckpt.checkpoint(self.down[i_level].block[i_block], hs[-1])  # type: ignore[index, operator]
-                    if len(self.down[i_level].attn) > 0:  # type: ignore[arg-type]
-                        h = ckpt.checkpoint(self.down[i_level].attn[i_block], h)  # type: ignore[index, operator]
+                    h = ckpt.checkpoint(self.down[i_level].block[i_block], hs[-1])
+                    if len(self.down[i_level].attn) > 0:
+                        h = ckpt.checkpoint(self.down[i_level].attn[i_block], h)
                 else:
-                    h = self.down[i_level].block[i_block](hs[-1])  # type: ignore[index, operator]
-                    if len(self.down[i_level].attn) > 0:  # type: ignore[arg-type]
-                        h = self.down[i_level].attn[i_block](h)  # type: ignore[index, operator]
+                    h = self.down[i_level].block[i_block](hs[-1])
+                    if len(self.down[i_level].attn) > 0:
+                        h = self.down[i_level].attn[i_block](h)
                 hs.append(h)
             if i_level != self.num_resolutions - 1:
                 if torch.is_grad_enabled() and self.gradient_checkpointing:
-                    hs.append(ckpt.checkpoint(self.down[i_level].downsample, hs[-1]))  # type: ignore[operator]
+                    hs.append(ckpt.checkpoint(self.down[i_level].downsample, hs[-1]))
                 else:
-                    hs.append(self.down[i_level].downsample(hs[-1]))  # type: ignore[operator]
+                    hs.append(self.down[i_level].downsample(hs[-1]))
 
         # middle
         h = hs[-1]
         if torch.is_grad_enabled() and self.gradient_checkpointing:
-            h = ckpt.checkpoint(self.mid.block_1, h)  # type: ignore[operator]
-            h = ckpt.checkpoint(self.mid.attn_1, h)  # type: ignore[operator]
-            h = ckpt.checkpoint(self.mid.block_2, h)  # type: ignore[operator]
+            h = ckpt.checkpoint(self.mid.block_1, h)
+            h = ckpt.checkpoint(self.mid.attn_1, h)
+            h = ckpt.checkpoint(self.mid.block_2, h)
         else:
-            h = self.mid.block_1(h)  # type: ignore[operator]
-            h = self.mid.attn_1(h)  # type: ignore[operator]
-            h = self.mid.block_2(h)  # type: ignore[operator]
+            h = self.mid.block_1(h)
+            h = self.mid.attn_1(h)
+            h = self.mid.block_2(h)
         # end
         h = self.norm_out(h)
         h = swish(h)
@@ -300,13 +322,13 @@ class Decoder(nn.Module):
 
         # middle
         if torch.is_grad_enabled() and self.gradient_checkpointing:
-            h = ckpt.checkpoint(self.mid.block_1, h)  # type: ignore[operator]
-            h = ckpt.checkpoint(self.mid.attn_1, h)  # type: ignore[operator]
-            h = ckpt.checkpoint(self.mid.block_2, h)  # type: ignore[operator]
+            h = ckpt.checkpoint(self.mid.block_1, h)
+            h = ckpt.checkpoint(self.mid.attn_1, h)
+            h = ckpt.checkpoint(self.mid.block_2, h)
         else:
-            h = self.mid.block_1(h)  # type: ignore[operator]
-            h = self.mid.attn_1(h)  # type: ignore[operator]
-            h = self.mid.block_2(h)  # type: ignore[operator]
+            h = self.mid.block_1(h)
+            h = self.mid.attn_1(h)
+            h = self.mid.block_2(h)
 
         # cast to proper dtype
         h = h.to(upscale_dtype)
@@ -314,18 +336,18 @@ class Decoder(nn.Module):
         for i_level in reversed(range(self.num_resolutions)):
             for i_block in range(self.num_res_blocks + 1):
                 if torch.is_grad_enabled() and self.gradient_checkpointing:
-                    h = ckpt.checkpoint(self.up[i_level].block[i_block], h)  # type: ignore[index, operator]
-                    if len(self.up[i_level].attn) > 0:  # type: ignore[arg-type]
-                        h = ckpt.checkpoint(self.up[i_level].attn[i_block], h)  # type: ignore[index, operator]
+                    h = ckpt.checkpoint(self.up[i_level].block[i_block], h)
+                    if len(self.up[i_level].attn) > 0:
+                        h = ckpt.checkpoint(self.up[i_level].attn[i_block], h)
                 else:
-                    h = self.up[i_level].block[i_block](h)  # type: ignore[index, operator]
-                    if len(self.up[i_level].attn) > 0:  # type: ignore[arg-type]
-                        h = self.up[i_level].attn[i_block](h)  # type: ignore[index, operator]
+                    h = self.up[i_level].block[i_block](h)
+                    if len(self.up[i_level].attn) > 0:
+                        h = self.up[i_level].attn[i_block](h)
             if i_level != 0:
                 if torch.is_grad_enabled() and self.gradient_checkpointing:
-                    h = ckpt.checkpoint(self.up[i_level].upsample, h)  # type: ignore[operator]
+                    h = ckpt.checkpoint(self.up[i_level].upsample, h)
                 else:
-                    h = self.up[i_level].upsample(h)  # type: ignore[operator]
+                    h = self.up[i_level].upsample(h)
 
         # end
         h = self.norm_out(h)
@@ -334,7 +356,14 @@ class Decoder(nn.Module):
         return h
 
 
-class AutoEncoder(nn.Module):
+class AutoEncoder(nn.Module, OstrisModelMixin):
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict):
+        # small decoder builds report 96ch at the first decoder block
+        if state_dict["decoder.up.0.block.0.conv1.bias"].shape[0] == 96:
+            return AutoEncoderSmallDecoderParams()
+        return AutoEncoderParams()
+
     def __init__(self, params: AutoEncoderParams):
         super().__init__()
         self.params = params
@@ -346,10 +375,13 @@ class AutoEncoder(nn.Module):
             num_res_blocks=params.num_res_blocks,
             z_channels=params.z_channels,
         )
+        decoder_ch = params.ch
+        if hasattr(params, "ch_encoder"):
+            decoder_ch = params.ch_encoder
         self.decoder = Decoder(
             resolution=params.resolution,
             in_channels=params.in_channels,
-            ch=params.ch,
+            ch=decoder_ch,
             out_ch=params.out_ch,
             ch_mult=params.ch_mult,
             num_res_blocks=params.num_res_blocks,
@@ -369,7 +401,7 @@ class AutoEncoder(nn.Module):
         self._gradient_checkpointing = False
 
     @property
-    def gradient_checkpointing(self) -> bool:
+    def gradient_checkpointing(self):
         return self._gradient_checkpointing
 
     @gradient_checkpointing.setter
@@ -378,18 +410,52 @@ class AutoEncoder(nn.Module):
         self.encoder.gradient_checkpointing = value
         self.decoder.gradient_checkpointing = value
 
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
     def enable_gradient_checkpointing(self):
         self.gradient_checkpointing = True
         self.encoder.enable_gradient_checkpointing()
         self.decoder.enable_gradient_checkpointing()
 
-    @property
-    def device(self) -> torch.device:
-        return next(self.parameters()).device
+    def normalize(self, z):
+        self.bn.eval()
+        return self.bn(z)
 
-    @property
-    def dtype(self) -> torch.dtype:
-        return next(self.parameters()).dtype
+    def inv_normalize(self, z):
+        self.bn.eval()
+        s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn_eps)
+        m = self.bn.running_mean.view(1, -1, 1, 1)
+        return z * s + m
+
+    def encode(self, x: Tensor) -> Tensor:
+        moments = self.encoder(x)
+        mean = torch.chunk(moments, 2, dim=1)[0]
+
+        z = rearrange(
+            mean,
+            "... c (i pi) (j pj)  -> ... (c pi pj) i j",
+            pi=self.ps[0],
+            pj=self.ps[1],
+        )
+        z = self.normalize(z)
+        return z
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = self.inv_normalize(z)
+        z = rearrange(
+            z,
+            "... (c pi pj) i j -> ... c (i pi) (j pj)",
+            pi=self.ps[0],
+            pj=self.ps[1],
+        )
+        dec = self.decoder(z)
+        return dec
 
 
 _NUM_RESOLUTIONS = 4

toolkit/models/v2/vae/ltx2.py

@@ -0,0 +1,16 @@
+from diffusers.models.autoencoders import (
+    AutoencoderKLLTX2Audio as DiffusersAutoencoderKLLTX2Audio,
+)
+from diffusers.models.autoencoders import (
+    AutoencoderKLLTX2Video as DiffusersAutoencoderKLLTX2Video,
+)
+
+from .._mixin import OstrisModelMixin
+
+
+class LTX2VideoVAE(DiffusersAutoencoderKLLTX2Video, OstrisModelMixin):
+    aitk_subfolder = "vae"
+
+
+class LTX2AudioVAE(DiffusersAutoencoderKLLTX2Audio, OstrisModelMixin):
+    aitk_subfolder = "audio_vae"

toolkit/models/v2/vae/qwen_image.py

@@ -0,0 +1,86 @@
+import torch
+from diffusers import AutoencoderKLQwenImage
+
+from .._mixin import OstrisModelMixin
+
+
+class QwenImageVAE(AutoencoderKLQwenImage, OstrisModelMixin):
+    """The Qwen-Image (wan-style video) VAE, shared by the qwen_image family,
+    nucleus_image and krea2."""
+
+    aitk_subfolder = "vae"
+
+
+class QwenImageVAEHolderMixin:
+    """BaseModel-side encode_images/decode_latents for models whose self.vae
+    is the Qwen-Image VAE: it is a video VAE, so images ride in a single-frame
+    dim and latents are normalized with the config's latents_mean/std."""
+
+    # tile the decode when low_vram (decode only; encode stays untiled)
+    vae_decode_tiled_on_low_vram = False
+
+    def encode_images(self, image_list, 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()
+        self.vae.requires_grad_(False)
+
+        image_list = [image.to(device, dtype=dtype) for image in image_list]
+        images = torch.stack(image_list).to(device, dtype=dtype)
+        images = images.unsqueeze(2)  # add the frame dim
+        latents = self.vae.encode(images).latent_dist.sample()
+
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
+            1, self.vae.config.z_dim, 1, 1, 1
+        ).to(latents.device, latents.dtype)
+
+        latents = (latents - latents_mean) * latents_std
+        latents = latents.squeeze(2)  # drop the frame dim
+        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(device, dtype=dtype)
+        latents = latents.unsqueeze(2)  # add the frame dim
+
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(self.vae.config.latents_std)
+            .view(1, self.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents = latents * latents_std + latents_mean
+
+        # full-resolution decode spikes VRAM; models opt in to tiling it
+        tiled = self.vae_decode_tiled_on_low_vram and self.model_config.low_vram
+        if tiled:
+            self.vae.enable_tiling()
+        try:
+            images = self.vae.decode(latents).sample
+        finally:
+            if tiled:
+                self.vae.disable_tiling()
+
+        images = images.squeeze(2)  # drop the frame dim
+        return images.to(device, dtype=dtype)

toolkit/models/v2/vae/wan.py

@@ -0,0 +1,9 @@
+from diffusers import AutoencoderKLWan
+
+from .._mixin import OstrisModelMixin
+
+
+class WanVAE(AutoencoderKLWan, OstrisModelMixin):
+    """The wan 2.1/2.2 causal video VAE."""
+
+    aitk_subfolder = "vae"