Move the offload and quantization to the base model

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:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing Text Encoder")
+        self.text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
-            quantize(self.text_encoder, weights=get_qtype(self.model_config.qtype_te))
+        flush()
-            freeze(self.text_encoder)
-            flush()
         
         self.vae = models["vae"]
         

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,13 +3,11 @@ 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
@@ -17,7 +15,6 @@ 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
@@ -287,62 +284,23 @@ class AnimaModel(BaseModel):
         transformer = pipe.transformer
         text_conditioner = pipe.text_conditioner
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing Transformer")
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
-            quantize_model(self, transformer)
-            flush()
 
-            self.print_and_status_update("Quantizing Text Conditioner")
+        # the text conditioner rides the transformer quantize flag (at qtype_te)
-            quantize(text_conditioner, weights=get_qtype(self.model_config.qtype_te))
+        # but takes the text-encoder offload/placement policy
-            freeze(text_conditioner)
+        tc_kwargs = self.component_load_kwargs("te")
-            flush()
+        tc_kwargs["qtype"] = (
+            self.model_config.qtype_te if self.model_config.quantize else None
+        )
+        text_conditioner.aitk_post_load(**tc_kwargs)
+        flush()
 
-        if (
+        # quantize + offload + placement, all driven by model_config
-            self.model_config.layer_offloading
+        pipe.text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
-            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)
         pipe.text_encoder.requires_grad_(False)
         pipe.text_encoder.eval()
-
+        flush()
-        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()
 
         self.noise_scheduler = pipe.scheduler
         self.vae = pipe.vae

extensions_built_in/diffusion_models/boogu_image/boogu_image.py

@@ -36,9 +36,9 @@ 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
@@ -164,16 +164,8 @@ class BooguImageModel(BaseModel):
         if attention_backend != "native":
             transformer.set_attention_backend(attention_backend)
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing transformer")
+        transformer.aitk_post_load(**self.component_load_kwargs("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)
         flush()
 
         # --- instruction encoder (Qwen3-VL) + processor ---
@@ -203,18 +195,8 @@ class BooguImageModel(BaseModel):
             )
         flush()
 
-        if self.model_config.quantize_te:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing instruction encoder")
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
-            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)
         flush()
 
         # --- VAE (FLUX AutoencoderKL) ---

extensions_built_in/diffusion_models/chroma/chroma_model.py

@@ -13,7 +13,6 @@ from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
 from toolkit.accelerator import unwrap_model
 from optimum.quanto import QTensor
-from toolkit.util.quantize import quantize, quantize_model
 from .pipeline import ChromaPipeline, prepare_latent_image_ids
 from einops import rearrange, repeat
 import random
@@ -158,27 +157,21 @@ class ChromaModel(BaseModel):
             transformer = Chroma.load_from_state_dict(load_file(model_path, "cpu"), dtype)
         # add dtype, not sure why it doesnt have it
         transformer.dtype = dtype
-        transformer.to(self.quantize_device, dtype=dtype)
 
         transformer.config = FakeConfig()
         transformer.config.num_layers = transformer.params.depth
         transformer.config.num_single_layers = transformer.params.depth_single_blocks
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            # block-streaming quantize (handles dequant-on-save patching,
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
-            # excludes, ARA, and quantize_kwargs)
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            transformer.to(self.device_torch)
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
 
         flush()
 
         self.print_and_status_update("Loading T5")
         tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
-        text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype)
+        text_encoder_2 = T5TextEncoder.load(
-        self.prepare_text_encoder(text_encoder_2, dtype=dtype)
+            extras_path, **self.component_load_kwargs("te")
+        )
 
         # self.print_and_status_update("Loading CLIP")
         text_encoder = FakeCLIP(device=self.device_torch)

extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py

@@ -12,7 +12,6 @@ from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
 from toolkit.accelerator import unwrap_model
 from optimum.quanto import QTensor
-from toolkit.util.quantize import quantize, quantize_model
 from .pipeline import ChromaPipeline, prepare_latent_image_ids
 from einops import rearrange, repeat
 import random
@@ -137,27 +136,21 @@ class ChromaRadianceModel(BaseModel):
             transformer = Chroma.load_model(model_path, dtype=dtype)
         # add dtype, not sure why it doesnt have it
         transformer.dtype = dtype
-        transformer.to(self.quantize_device, dtype=dtype)
 
         transformer.config = FakeConfig()
         transformer.config.num_layers = transformer.params.depth
         transformer.config.num_single_layers = transformer.params.depth_single_blocks
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            # block-streaming quantize (handles dequant-on-save patching,
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
-            # excludes, ARA, and quantize_kwargs)
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            transformer.to(self.device_torch)
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
 
         flush()
 
         self.print_and_status_update("Loading T5")
         tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
-        text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype)
+        text_encoder_2 = T5TextEncoder.load(
-        self.prepare_text_encoder(text_encoder_2, dtype=dtype)
+            extras_path, **self.component_load_kwargs("te")
+        )
 
         # self.print_and_status_update("Loading CLIP")
         text_encoder = FakeCLIP(device=self.device_torch)

extensions_built_in/diffusion_models/ernie_image/ernie_image.py

@@ -13,9 +13,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, quantize_model
-from toolkit.memory_management import MemoryManager
 
 from transformers import AutoTokenizer, AutoModel
 
@@ -88,29 +85,10 @@ class ErnieImageModel(BaseModel):
             if os.path.exists(te_folder_path):
                 base_model_path = model_path
 
-        transformer = ErnieImageTransformer2DModel.load_model(model_path, dtype=dtype)
+        # load + quantize + offload + placement, all driven by model_config
-
+        transformer = ErnieImageTransformer2DModel.load(
-        if self.model_config.quantize:
+            model_path, **self.component_load_kwargs("transformer")
-            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()
 
@@ -118,29 +96,11 @@ class ErnieImageModel(BaseModel):
         tokenizer = AutoTokenizer.from_pretrained(
             base_model_path, subfolder="tokenizer", torch_dtype=dtype
         )
-        text_encoder = Mistral3ModelEncoder.load_model(
+        text_encoder = Mistral3ModelEncoder.load(
-            base_model_path, subfolder="text_encoder", dtype=dtype
+            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 = Flux2KLVAE.load_model(            base_model_path, dtype=dtype
         ).to(self.device_torch, dtype=dtype)

extensions_built_in/diffusion_models/ernie_image/transformer.py

@@ -298,6 +298,9 @@ class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin, OstrisModelMixin):
     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 + offload + placement, all driven by model_config:
-            # quantize_model handles qtype selection, exclusions and device
+        # component_load_kwargs derives qtype (incl. an accuracy recovery
-            # juggling, and leaves the model on CPU
+        # adapter), the layer-offload fraction and the target device
-            self.print_and_status_update("Quantizing transformer")
+        # (low_vram parks on CPU); aitk_post_load applies them. Models whose
-            quantize_model(self, transformer)
+        # checkpoint sourcing is standard can collapse the build + this into
-            flush()
+        # one call: ExampleTransformer2DModel.load(path, **kwargs).
-
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
-        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)
         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")

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):
+class ExampleTransformer2DModel(nn.Module, OstrisModelMixin):
-    """The denoiser. Plain ``nn.Module`` on purpose.
+    """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
+        # what get_transformer_block_names() returns, which the LoRA code uses
-        # LoRA code uses for block targeting / "transformer only" training.
+        # 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

@@ -13,7 +13,7 @@ from toolkit.prompt_utils import PromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
 from toolkit.accelerator import unwrap_model
 from optimum.quanto import QTensor
-from toolkit.util.quantize import quantize, quantize_model
+
 from .src import FLitePipeline, DiT
 
 if TYPE_CHECKING:
@@ -76,27 +76,17 @@ class FLiteModel(BaseModel):
 
         self.print_and_status_update("Loading transformer")
 
-        transformer = DiT.load_model(model_path, dtype=dtype)
+        transformer = DiT.load(
-        
+            model_path, **self.component_load_kwargs("transformer")
-        transformer.to(self.quantize_device, dtype=dtype)
+        )
-
-        if self.model_config.quantize:
-            # block-streaming quantize (handles dequant-on-save patching,
-            # excludes, ARA, and quantize_kwargs)
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            transformer.to(self.device_torch)
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
 
         flush()
 
         self.print_and_status_update("Loading T5")
         tokenizer = T5TextEncoder.load_tokenizer(extras_path, subfolder="tokenizer")
-        text_encoder = T5TextEncoder.load_model(
+        text_encoder = T5TextEncoder.load(
-            extras_path, dtype=dtype, subfolder="text_encoder"
+            extras_path, subfolder="text_encoder", **self.component_load_kwargs("te")
         )
-        self.prepare_text_encoder(text_encoder, dtype=dtype)
 
         self.noise_scheduler = FLiteModel.get_train_scheduler()
         

extensions_built_in/diffusion_models/flux2/flux2_klein_model.py

@@ -1,10 +1,7 @@
 from .flux2_model import Flux2Model
 from transformers import Qwen3ForCausalLM, Qwen2Tokenizer
 from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
-from optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype
 from toolkit.config_modules import ModelConfig
-from toolkit.memory_management.manager import MemoryManager
 from toolkit.basic import flush
 from .src.model import Klein9BParams, Klein4BParams
 
@@ -41,27 +38,11 @@ class Flux2KleinModel(Flux2Model):
         dtype = self.torch_dtype
         self.print_and_status_update("Loading Qwen3")
 
-        text_encoder = Qwen3TextEncoder.load_model(
+        # load + quantize + offload + placement, all driven by model_config
-            self.flux2_klein_te_path, dtype=dtype, subfolder=""
+        text_encoder = Qwen3TextEncoder.load(
+            self.flux2_klein_te_path, subfolder="", **self.component_load_kwargs("te")
         )
-        if self.model_config.quantize_te:
+        flush()
-            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,
-            )
 
         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,10 +12,8 @@ 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 optimum.quanto import QTensor
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
 
 from transformers import AutoProcessor, Mistral3ForConditionalGeneration
 from toolkit.models.v2.text_encoders.mistral3 import Mistral3TextEncoder
@@ -100,29 +97,12 @@ class Flux2Model(BaseModel):
         dtype = self.torch_dtype
         self.print_and_status_update("Loading Mistral")
 
-        text_encoder = Mistral3TextEncoder.load_model(
+        # load + quantize + offload + placement, all driven by model_config
-            MISTRAL_PATH, dtype=dtype, subfolder=""
+        text_encoder = Mistral3TextEncoder.load(
+            MISTRAL_PATH, subfolder="", **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_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,
-            )
-
         tokenizer = AutoProcessor.from_pretrained(MISTRAL_PATH)
         return text_encoder, tokenizer
 
@@ -151,32 +131,10 @@ class Flux2Model(BaseModel):
             transformer_state_dict, dtype, config=self.get_flux2_params()
         )
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            # patch the state dict method
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
-            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)
         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")

extensions_built_in/diffusion_models/flux2/src/model.py

@@ -57,6 +57,10 @@ class FakeConfig:
 
 
 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/flux_kontext/flux_kontext.py

@@ -20,7 +20,7 @@ from toolkit.models.flux import add_model_gpu_splitter_to_flux, bypass_flux_guid
 from toolkit.accelerator import get_accelerator, unwrap_model
 from optimum.quanto import QTensor
 from toolkit.util.mask import generate_random_mask, random_dialate_mask
-from toolkit.util.quantize import quantize, quantize_model
+
 from einops import rearrange, repeat
 import random
 import torch.nn.functional as F
@@ -38,8 +38,6 @@ scheduler_config = {
     "use_dynamic_shifting": True
 }
 
-
-
 class FluxKontextModel(BaseModel):
     arch = "flux_kontext"
 
@@ -93,24 +91,17 @@ class FluxKontextModel(BaseModel):
                 base_model_path = model_path
 
         self.print_and_status_update("Loading transformer")
-        transformer = FluxTransformer2DModel.load_model(model_path, dtype=dtype)
+        transformer = FluxTransformer2DModel.load(
-        transformer.to(self.quantize_device, dtype=dtype)
+            model_path, **self.component_load_kwargs("transformer")
-
+        )
-        if self.model_config.quantize:
-            # block-streaming quantize (handles dequant-on-save patching,
-            # excludes, ARA, and quantize_kwargs)
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            transformer.to(self.device_torch)
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
 
         flush()
 
         self.print_and_status_update("Loading T5")
         tokenizer_2 = T5TextEncoder.load_tokenizer(base_model_path)
-        text_encoder_2 = T5TextEncoder.load_model(base_model_path, dtype=dtype)
+        text_encoder_2 = T5TextEncoder.load(
-        self.prepare_text_encoder(text_encoder_2, dtype=dtype)
+            base_model_path, **self.component_load_kwargs("te")
+        )
 
         self.print_and_status_update("Loading CLIP")
         text_encoder = CLIPTextEncoder.load_model(

extensions_built_in/diffusion_models/hidream/hidream_model.py

@@ -22,7 +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 .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
@@ -97,62 +96,23 @@ class HidreamModel(BaseModel):
             use_fast=False
         )
         
-        text_encoder_4 = LlamaTextEncoder.load_model(
+        # load + quantize + offload + placement, all driven by model_config
+        text_encoder_4 = LlamaTextEncoder.load(
             llama_model_path,
-            dtype=torch.bfloat16,
             subfolder="",
             output_hidden_states=True,
             output_attentions=True,
+            **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.load_model(
+        transformer = self.hidream_transformer_class.load(
-            model_path, dtype=torch.bfloat16
+            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")
@@ -183,16 +143,11 @@ class HidreamModel(BaseModel):
         flush()
         self.print_and_status_update("Loading T5 encoders")
         
-        text_encoder_3 = T5TextEncoder.load_model(
+        # load + quantize + offload + placement, all driven by model_config
-            extras_path, dtype=torch.bfloat16, subfolder="text_encoder_3"
+        text_encoder_3 = T5TextEncoder.load(
-        ).to(self.device_torch, dtype=dtype)
+            extras_path, subfolder="text_encoder_3", **self.component_load_kwargs("te")
-        
+        )
-        if self.model_config.quantize_te:
+        flush()
-            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 = T5TextEncoder.load_tokenizer(
             extras_path, subfolder="tokenizer_3", use_fast=False

extensions_built_in/diffusion_models/hidream/hidream_o1_model.py

@@ -16,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
@@ -203,25 +201,8 @@ class HidreamO1Model(BaseModel):
                 torch_dtype=self.torch_dtype,
             )
         flush()
-        if not self.model_config.low_vram:
+        # quantize + offload + placement, all driven by model_config
-            transformer.to(self.device_torch)
+        transformer.aitk_post_load(**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=[],
-            )
-
         flush()
 
         # move over to device now if low vram

extensions_built_in/diffusion_models/ideogram4/ideogram4.py

@@ -17,9 +17,7 @@ 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 optimum.quanto import QTensor
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
-from optimum.quanto import freeze, QTensor
 
 import huggingface_hub
 from huggingface_hub.errors import EntryNotFoundError
@@ -365,58 +363,13 @@ class Ideogram4Model(BaseModel):
 
         transformer = self._load_transformer(base)
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing Transformer")
+        transformer.aitk_post_load(**self.component_load_kwargs("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)
         flush()
 
         tokenizer, text_encoder = self._load_text_encoder(base)
-        if self.model_config.quantize_te:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing Text Encoder")
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
-            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)
         flush()
 
         vae = self._load_vae(base)

extensions_built_in/diffusion_models/ideogram4/src/transformer.py

@@ -360,6 +360,17 @@ class Ideogram4FinalLayer(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

@@ -47,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 (
@@ -383,55 +382,12 @@ class Krea2Model(QwenImageVAEHolderMixin, BaseModel):
             if self.model_config.qtype == "qfloat8":
                 self.model_config.qtype = "float8"
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing transformer")
+        transformer.aitk_post_load(**self.component_load_kwargs("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)
         flush()
 
         tokenizer, processor, vl_processor, text_encoder = self._load_text_encoder()
-        if self.model_config.quantize_te:
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
-            self.print_and_status_update("Quantizing text encoder")
-            text_encoder.to(self.device_torch)
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
-        if (
-            self.model_config.layer_offloading
-            and self.model_config.layer_offloading_text_encoder_percent > 0
-        ):
-            MemoryManager.attach(
-                text_encoder,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_text_encoder_percent,
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving text encoder to CPU")
-            text_encoder.to("cpu")
-        else:
-            text_encoder.to(self.device_torch)
         flush()
 
         vae = self._load_vae()

extensions_built_in/diffusion_models/krea2/src/mmdit.py

@@ -411,6 +411,14 @@ class SingleStreamBlock(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/ltx2.py

@@ -17,9 +17,6 @@ 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
@@ -307,33 +304,8 @@ class LTX2Model(BaseModel):
 
             transformer = LTX2VideoTransformer3DModel.load_model(model_path, dtype=dtype)
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing Transformer")
+        transformer.aitk_post_load(**self.component_load_kwargs("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")
 
         flush()
 
@@ -441,25 +413,8 @@ class LTX2Model(BaseModel):
         text_encoder.model.vision_tower = None
         flush()
         
-        if self.model_config.quantize_te:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing Text Encoder")
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
-            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
-                ],
-            )
-
         text_encoder.to(self.device_torch, dtype=dtype)
         flush()
 
@@ -1375,38 +1330,11 @@ class LTX25Model(LTX2Model):
             pin_stored_fp32(transformer)
         flush()
 
-        if self.model_config.quantize:
+        if num_quantized_dit:
-            if num_quantized_dit:
+            # pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize
-                self.print_and_status_update(
+            transformer.aitk_is_quantized = True
-                    "Transformer is pre-quantized (ConvRot); skipping quantize"
+        # quantize + offload + placement, all driven by model_config
-                )
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
-            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")
         flush()
 
         self.print_and_status_update("Loading connectors")
@@ -1433,32 +1361,11 @@ class LTX25Model(LTX2Model):
         del te_state_dict
         flush()
 
-        if self.model_config.quantize_te:
+        if num_quantized_te:
-            if num_quantized_te:
+            # pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize
-                self.print_and_status_update(
+            text_encoder.aitk_is_quantized = True
-                    "Text encoder is pre-quantized (ConvRot); skipping quantize"
+        # quantize + offload + placement, all driven by model_config
-                )
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
-            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,
-            )
 
         text_encoder.to(self.device_torch)
         text_encoder.requires_grad_(False)

extensions_built_in/diffusion_models/mageflow/mageflow.py

@@ -36,7 +36,6 @@ from safetensors.torch import load_file, save_file
 import huggingface_hub
 from transformers import AutoProcessor, AutoTokenizer
 from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
-from optimum.quanto import freeze
 
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.models.base_model import BaseModel
@@ -47,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
@@ -262,50 +259,13 @@ class MageFlowModel(BaseModel):
 
         transformer = self._load_transformer()
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing transformer")
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
-            quantize_model(self, transformer)
-            flush()
-
-        if (
-            self.model_config.layer_offloading
-            and self.model_config.layer_offloading_transformer_percent > 0
-        ):
-            MemoryManager.attach(
-                transformer,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_transformer_percent,
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
         flush()
 
         tokenizer, vl_processor, text_encoder = self._load_text_encoder()
-        if self.model_config.quantize_te:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing text encoder")
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
-            text_encoder.to(self.device_torch)
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
-        if (
-            self.model_config.layer_offloading
-            and self.model_config.layer_offloading_text_encoder_percent > 0
-        ):
-            MemoryManager.attach(
-                text_encoder,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_text_encoder_percent,
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving text encoder to CPU")
-            text_encoder.to("cpu")
-        else:
-            text_encoder.to(self.device_torch)
         flush()
 
         vae = self._load_vae()

extensions_built_in/diffusion_models/mageflow/src/transformer.py

@@ -684,6 +684,26 @@ class MageFlowParams:
 
 
 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/minimax_h3/minimax_h3.py

@@ -49,7 +49,6 @@ 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
@@ -64,8 +63,6 @@ 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
 
@@ -498,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:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing transformer")
+        transformer.aitk_post_load(**self.component_load_kwargs("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)
         flush()
 
         tokenizer, processor, text_encoder = self._load_text_encoder()
-        te_prequantized = any(
+        if any(isinstance(m, OstrisLinear) for m in text_encoder.modules()):
-            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
-        if self.model_config.quantize_te and not te_prequantized:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing text encoder")
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
-            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)
         flush()
 
         vae_bundle = self._load_vaes()

extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py

@@ -14,8 +14,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.accelerator import unwrap_model
-from toolkit.util.quantize import quantize_model
-from toolkit.memory_management import MemoryManager
 
 from transformers import Qwen3VLProcessor
 import torch.nn.functional as F
@@ -100,26 +98,8 @@ class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel):
                 if isinstance(m, SwiGLUExperts):
                     m.use_grouped_mm = False
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing Transformer")
+        transformer.aitk_post_load(**self.component_load_kwargs("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")
 
         flush()
 
@@ -127,8 +107,9 @@ class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel):
         tokenizer = Qwen3VLProcessor.from_pretrained(
             base_model_path, subfolder="processor", torch_dtype=dtype
         )
-        text_encoder = Qwen3VLTextEncoder.load_model(base_model_path, dtype=dtype)
+        text_encoder = Qwen3VLTextEncoder.load(
-        self.prepare_text_encoder(text_encoder, dtype=dtype)
+            base_model_path, **self.component_load_kwargs("te")
+        )
 
         self.print_and_status_update("Loading VAE")
         vae = QwenImageVAE.load_model(

extensions_built_in/diffusion_models/omnigen2/__init__.py

@@ -14,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, quantize_model
 from .src.pipelines.omnigen2.pipeline_omnigen2 import OmniGen2Pipeline
 from .src.models.transformers import OmniGen2Transformer2DModel
 from .src.models.transformers.repo import OmniGen2RotaryPosEmbed
@@ -80,41 +78,19 @@ class OmniGen2Model(BaseModel):
             extras_path, subfolder="processor", use_fast=True
         )
 
-        mllm = Qwen25VLTextEncoder.load_model(
+        # load + quantize + offload + placement, all driven by model_config
-            extras_path, dtype=torch.bfloat16, subfolder="mllm"
+        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.load_model(
+        transformer = OmniGen2Transformer2DModel.load(
-            model_path, dtype=torch.bfloat16
+            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")
-            quantize_model(self, transformer)
-            if not self.low_vram:
-                transformer.to(self.device_torch)
-
-        if self.low_vram:
-            # unload it for now
-            transformer.to("cpu")
-
         flush()
 
         self.print_and_status_update("Loading vae")

extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py

@@ -43,7 +43,7 @@ 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
@@ -56,7 +56,7 @@ 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
@@ -127,36 +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.load_model(model_path, dtype=dtype)
+        # load + quantize + offload + placement, all driven by model_config
-        transformer.to(dtype=dtype)
+        transformer = PRXTransformer2DModel.load(
-        flush()
+            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.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 = Qwen3VLTextOnlyEncoder.load_model(model_path, dtype=dtype)
+        text_encoder = Qwen3VLTextOnlyEncoder.load(
-        text_encoder.to(self.te_device_torch)
+            model_path, **self.component_load_kwargs("te")
+        )
         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)

extensions_built_in/diffusion_models/qwen_image/qwen_image.py

@@ -14,7 +14,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.accelerator import get_accelerator, unwrap_model
-from toolkit.util.quantize import quantize_model
 import torch.nn.functional as F
 from toolkit.memory_management import MemoryManager
 from toolkit.metadata import get_meta_for_safetensors
@@ -102,35 +101,10 @@ class QwenImageModel(QwenImageVAEHolderMixin, BaseModel):
             if os.path.exists(te_folder_path):
                 base_model_path = model_path
 
-        transformer = QwenImageTransformer2DModel.load_model(
+        # load + quantize + offload + placement, all driven by model_config
-            model_path,
+        transformer = QwenImageTransformer2DModel.load(
-            dtype=model_dtype,
+            model_path, **self.component_load_kwargs("transformer")
-            use_comfy_weights=self.model_config.model_kwargs.get(
-                "use_comfy_weights", True
-            ),
         )
-
-        if self.model_config.quantize and not getattr(
-            transformer, "aitk_is_quantized", False
-        ):
-            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")
-
         flush()
 
         self.print_and_status_update("Text Encoder")
@@ -138,11 +112,12 @@ class QwenImageModel(QwenImageVAEHolderMixin, BaseModel):
         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.drop_vision_tower()
 
-        self.prepare_text_encoder(text_encoder, dtype=dtype)
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
 
         self.print_and_status_update("Loading VAE")
         vae = QwenImageVAE.load_model(base_model_path, 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

@@ -286,62 +286,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.load(
-        transformer_1 = WanTransformer3DModel.load_model(
+            transformer_path_1, subfolder=subfolder_1, **per_kwargs
-            transformer_path_1, dtype=dtype, subfolder=subfolder_1
         )
-
         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.load(
-        transformer_2 = WanTransformer3DModel.load_model(
+            transformer_path_2, subfolder=subfolder_2, **per_kwargs
-            transformer_path_2, dtype=dtype, subfolder=subfolder_2
         )
-
         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")
@@ -362,18 +325,13 @@ class Wan2214bModel(Wan21):
             
         
         if layer_offloading_transformer:
-            MemoryManager.attach(
+            for t in (transformer_1, transformer_2):
-                transformer_1,
+                MemoryManager.attach(
-                self.device_torch,
+                    t,
-                offload_percent=self.model_config.layer_offloading_transformer_percent,
+                    self.device_torch,
-                ignore_modules=[transformer_1.scale_shift_table] + [block.scale_shift_table for block in transformer_1.blocks]
+                    offload_percent=self.model_config.layer_offloading_transformer_percent,
-            )
+                    ignore_modules=t.get_offload_ignore_modules(),
-            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]
-            )
 
         return transformer
 

extensions_built_in/diffusion_models/z_image/z_image.py

@@ -12,7 +12,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.accelerator import unwrap_model
-from toolkit.util.quantize import quantize_model
 from toolkit.memory_management import MemoryManager
 from toolkit.metadata import get_meta_for_safetensors
 from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
@@ -182,19 +181,9 @@ 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
@@ -223,37 +212,15 @@ 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
+        # quantize + offload + placement, all driven by model_config
-        # into full precision weights first (or the checkpoint was pre-quantized)
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
-        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")
-
         flush()
 
         self.print_and_status_update("Text Encoder")
         tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
-        text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
+        text_encoder = Qwen3TextEncoder.load(
-        self.prepare_text_encoder(text_encoder, dtype=dtype)
+            base_model_path, **self.component_load_kwargs("te")
+        )
 
         self.print_and_status_update("Loading VAE")
         vae = KLVAE.load_model(base_model_path, dtype=dtype)

extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py

@@ -9,8 +9,6 @@ import torch.nn.functional as F
 import yaml
 from toolkit.basic import flush
 from toolkit.accelerator import unwrap_model
-from toolkit.util.quantize import quantize_model
-from toolkit.memory_management import MemoryManager
 
 from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
 from toolkit.models.FakeVAE import FakeVAE
@@ -441,35 +439,16 @@ class ZImageL2PModel(ZImageModel):
             if self.model_config.qtype == "qfloat8":
                 self.model_config.qtype = "float8"
 
-        if self.model_config.quantize:
+        # quantize + offload + placement, all driven by model_config
-            self.print_and_status_update("Quantizing Transformer")
+        transformer.aitk_post_load(**self.component_load_kwargs("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")
 
         flush()
 
         self.print_and_status_update("Text Encoder")
         tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
-        text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
+        text_encoder = Qwen3TextEncoder.load(
-        self.prepare_text_encoder(text_encoder, dtype=dtype)
+            base_model_path, **self.component_load_kwargs("te")
+        )
 
         self.print_and_status_update("Loading VAE")
         # vae = AutoencoderKL.from_pretrained(

extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py

@@ -13,8 +13,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.accelerator import unwrap_model
-from toolkit.util.quantize import 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
@@ -108,37 +106,16 @@ class ZetaChromaModel(BaseModel):
         )
         del transformer_state_dict
 
-        transformer.to(self.quantize_device, dtype=dtype)
+        # quantize + offload + placement, all driven by model_config
-
+        transformer.aitk_post_load(**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_pad_token,
-                    transformer.cap_pad_token,
-                ],
-            )
-
-        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 = Qwen3TextEncoder.load_tokenizer(base_model_path)
-        text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
+        text_encoder = Qwen3TextEncoder.load(
-        self.prepare_text_encoder(text_encoder, dtype=dtype)
+            base_model_path, **self.component_load_kwargs("te")
+        )
 
         self.print_and_status_update("Loading VAE")
         vae = FakeVAE(scaling_factor=1.0)

extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py

@@ -452,6 +452,13 @@ class SimpleMLPAdaLN(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()

toolkit/models/base_model.py

@@ -1628,37 +1628,41 @@ class BaseModel:
         else:
             self.text_encoder.to(*args, **kwargs)
 
-    def prepare_text_encoder(self, text_encoder, dtype=None):
+    def component_load_kwargs(self, role: str = "transformer", dtype=None):
-        """Standard post-load text-encoder policy: layer offloading, device
+        """kwargs for a v2 module's .load()/.aitk_post_load(), derived from
-        placement, then quantize_te. Skips quantization when the checkpoint
+        model_config: qtype (with the accuracy recovery adapter recombined),
-        loaded pre-quantized."""
+        offload fraction, devices, low_vram placement. Roles: "transformer",
-        from optimum.quanto import freeze
+        "te", "vae"."""
-        from toolkit.memory_management import MemoryManager
+        mc = self.model_config
-        from toolkit.util.quantize import get_qtype, quantize
+        qtype, offload = None, 0.0
-
+        if role == "transformer":
-        dtype = dtype if dtype is not None else self.torch_dtype
+            if mc.quantize:
-        if (
+                qtype = mc.qtype
-            self.model_config.layer_offloading
+                if mc.accuracy_recovery_adapter and "|" not in (qtype or ""):
-            and self.model_config.layer_offloading_text_encoder_percent > 0
+                    qtype = f"{qtype}|{mc.accuracy_recovery_adapter}"
-        ):
+            if mc.layer_offloading:
-            MemoryManager.attach(
+                offload = mc.layer_offloading_transformer_percent
-                text_encoder,
+        elif role == "te":
-                self.device_torch,
+            if mc.quantize_te:
-                offload_percent=self.model_config.layer_offloading_text_encoder_percent,
+                qtype = mc.qtype_te
-            )
+            if mc.layer_offloading:
-
+                offload = mc.layer_offloading_text_encoder_percent
-        text_encoder.to(self.device_torch, dtype=dtype)
+        if dtype is None:
-        flush()
+            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 self.model_config.quantize_te and not getattr(
+        if mc.low_vram and role in ("transformer", "te"):
-            text_encoder, "aitk_is_quantized", False
+            device = "cpu"
-        ):
+        elif role == "vae":
-            self.print_and_status_update("Quantizing Text Encoder")
+            device = self.vae_device_torch
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
+        return dict(
-            freeze(text_encoder)
+            qtype=qtype,
-            flush()
+            offload=offload,
-        return text_encoder
+            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

toolkit/models/v2/PLANNING.md

@@ -408,6 +408,40 @@ 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.
+
 ## Testing
 
 - [x] `testing/test_model_loading.py`: per-arch load + one small sample through

toolkit/models/v2/_mixin.py

@@ -122,6 +122,11 @@ 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.
 
@@ -152,6 +157,131 @@ class OstrisModelMixin:
     # 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 []
+            )
+            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)
+            else:
+                status_fn(f"Quantizing ({qtype})")
+                quantize_module(
+                    self,
+                    qtype,
+                    device=quantize_device if quantize_device is not None else 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,
+                )
+            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,

toolkit/models/v2/diffusion_models/ltx2.py

@@ -21,6 +21,18 @@ class LTX2VideoTransformer3DModel(
     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"

toolkit/models/v2/diffusion_models/wan.py

@@ -80,6 +80,12 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin):
     def get_transformer_block_names(cls):
         return ["blocks"]
 
+    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

toolkit/models/v2/diffusion_models/z_image.py

@@ -26,6 +26,9 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
     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

toolkit/models/v2/text_encoders/gemma3.py

@@ -14,6 +14,9 @@ class Gemma3TextEncoder(Gemma3ForConditionalGeneration, OstrisTransformersMixin)
         # 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
@@ -28,5 +31,12 @@ try:
         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/wan21/wan21.py

@@ -6,7 +6,6 @@ from toolkit.accelerator import unwrap_model
 from toolkit.basic import flush
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.dequantize import patch_dequantization_on_save
-from toolkit.memory_management.manager import MemoryManager
 from toolkit.models.base_model import BaseModel
 from toolkit.prompt_utils import PromptEmbeds
 from transformers import AutoTokenizer, UMT5EncoderModel
@@ -42,7 +41,6 @@ from diffusers.pipelines.wan.pipeline_wan import XLA_AVAILABLE
 from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
 from typing import Any, Callable, Dict, List, Optional, Union
 from toolkit.models.wan21.wan_lora_convert import convert_to_diffusers, convert_to_original
-from toolkit.util.quantize import quantize_model
 from toolkit.models.v2.text_encoders.umt5 import UMT5TextEncoder
 from toolkit.models.v2.vae.wan import WanVAE
 from toolkit.metadata import get_meta_for_safetensors
@@ -344,23 +342,11 @@ class Wan21(BaseModel):
     
     def load_wan_transformer(self, transformer_path, subfolder=None):
         self.print_and_status_update("Loading transformer")
-        dtype = self.torch_dtype
-        transformer = WanTransformer3DModel.load_model(
-            transformer_path, dtype=dtype, subfolder=subfolder
-        )
 
         if self.model_config.split_model_over_gpus:
             raise ValueError(
                 "Splitting model over gpus is not supported for Wan2.1 models")
 
-        if self.model_config.low_vram:
-            # quantize on the device
-            transformer.to('cpu', dtype=dtype)
-            flush()
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
-            flush()
-
         if self.model_config.assistant_lora_path is not None or self.model_config.inference_lora_path is not None:
             raise ValueError(
                 "Assistant LoRA is not supported for Wan2.1 models currently")
@@ -369,23 +355,13 @@ class Wan21(BaseModel):
             raise ValueError(
                 "Loading LoRA is not supported for Wan2.1 models currently")
 
+        # load + quantize + offload + placement, all driven by model_config
+        transformer = WanTransformer3DModel.load(
+            transformer_path,
+            subfolder=subfolder,
+            **self.component_load_kwargs("transformer"),
+        )
         flush()
-        
-        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')
 
         return transformer
 
@@ -418,8 +394,9 @@ class Wan21(BaseModel):
         self.print_and_status_update("Loading UMT5EncoderModel")
         
         tokenizer = UMT5TextEncoder.load_tokenizer(te_path)
-        text_encoder = UMT5TextEncoder.load_model(te_path, dtype=dtype)
+        text_encoder = UMT5TextEncoder.load(
-        self.prepare_text_encoder(text_encoder, dtype=dtype)
+            te_path, **self.component_load_kwargs("te")
+        )
 
         if self.model_config.low_vram:
             print("Moving transformer back to GPU")

toolkit/util/quantize.py

@@ -261,6 +261,68 @@ def _has_quantizable_linear(module: torch.nn.Module, weights, exclude=None) -> b
     return False
 
 
+@torch.no_grad()
+def quantize_module(
+    module: torch.nn.Module,
+    qtype: str,
+    device=None,
+    dtype: torch.dtype = torch.bfloat16,
+    block_names: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
+    quantize_kwargs: Optional[dict] = None,
+    status_fn=print_acc,
+):
+    """Module-centric quantization: block-streamed (each repeated block moves
+    to ``device`` for the math and returns to cpu) with a whole-module pass
+    for the extras. This is the core the per-model loaders call; holders'
+    quantize_model wraps it with model_config plumbing."""
+    from toolkit.dequantize import patch_dequantization_on_save
+
+    patch_dequantization_on_save(module)
+    quantization_type = get_qtype(qtype)
+    exclude = list(exclude or [])
+    quantize_kwargs = quantize_kwargs or {}
+
+    all_blocks: List[torch.nn.Module] = []
+    for name in block_names or []:
+        # name may be a dotted path for models that nest their blocks
+        block_list = module
+        for part in name.split("."):
+            block_list = getattr(block_list, part, None)
+            if block_list is None:
+                break
+        if block_list is not None:
+            all_blocks += list(block_list)
+    if all_blocks:
+        status_fn(f" - quantizing {len(all_blocks)} blocks")
+    already_quantized = 0
+    for block in tqdm(all_blocks):
+        if not _has_quantizable_linear(block, quantization_type, exclude):
+            # pre-quantized checkpoint with a matching qtype: nothing in this
+            # block would change — skip the device round-trip and the dtype
+            # cast entirely so the load stays byte-identical
+            already_quantized += 1
+            continue
+        if device is not None:
+            block.to(device, dtype=dtype, non_blocking=True)
+        quantize(block, weights=quantization_type, exclude=exclude, **quantize_kwargs)
+        freeze(block)
+        # NOT non_blocking: an async D2H allocates the cpu destination in pinned
+        # memory, which the caching host allocator keeps forever — that silently
+        # retained a model-sized chunk of host ram
+        if device is not None:
+            block.to("cpu")
+    if already_quantized:
+        status_fn(
+            f" - {already_quantized} blocks already quantized with a matching qtype; left untouched"
+        )
+
+    status_fn(" - quantizing extras")
+    quantize(module, weights=quantization_type, exclude=exclude, **quantize_kwargs)
+    freeze(module)
+    return module
+
+
 def quantize_model(
     base_model: "BaseModel",
     model_to_quantize: torch.nn.Module,
@@ -279,153 +341,176 @@ def quantize_model(
     exclude_modules = base_model.get_quantization_exclude_modules() or []
 
     if base_model.model_config.accuracy_recovery_adapter is not None:
-        from toolkit.config_modules import NetworkConfig
+        attach_ara_and_quantize(
-        from toolkit.lora_special import LoRASpecialNetwork
+            base_model,
-
-        # we need to load and quantize with an accuracy recovery adapter
-        # todo handle hf repos
-        load_lora_path = base_model.model_config.accuracy_recovery_adapter
-
-        if not os.path.exists(load_lora_path):
-            # not local file, grab from the hub
-
-            path_split = load_lora_path.split("/")
-            if len(path_split) > 3:
-                raise ValueError(
-                    "The accuracy recovery adapter path must be a local path or for a hf repo, 'username/repo_name/filename.safetensors'."
-                )
-            repo_id = f"{path_split[0]}/{path_split[1]}"
-            print_acc(f"Grabbing lora from the hub: {load_lora_path}")
-            new_lora_path = hf_hub_download(
-                repo_id,
-                filename=path_split[-1],
-            )
-            # replace the path
-            load_lora_path = new_lora_path
-
-        # build the lora config based on the lora weights
-        lora_state_dict = load_file(load_lora_path)
-        
-        if hasattr(base_model, "convert_lora_weights_before_load"):
-            lora_state_dict = base_model.convert_lora_weights_before_load(lora_state_dict)
-        
-        network_config = {
-            "type": "lora",
-            "network_kwargs": {"only_if_contains": []},
-            "transformer_only": False,
-        }
-        first_key = list(lora_state_dict.keys())[0]
-        first_weight = lora_state_dict[first_key]
-        # if it starts with lycoris and includes lokr
-        if first_key.startswith("lycoris") and any(
-            "lokr" in key for key in lora_state_dict.keys()
-        ):
-            network_config["type"] = "lokr"
-        
-        network_kwargs = {}
-
-        # find firse loraA weight
-        if network_config["type"] == "lora":
-            linear_dim = None
-            for key, value in lora_state_dict.items():
-                if "lora_A" in key:
-                    linear_dim = int(value.shape[0])
-                    break
-            linear_alpha = linear_dim
-            network_config["linear"] = linear_dim
-            network_config["linear_alpha"] = linear_alpha
-
-            # we build the keys to match every key
-            only_if_contains = []
-            for key in lora_state_dict.keys():
-                contains_key = key.split(".lora_")[0]
-                if contains_key not in only_if_contains:
-                    only_if_contains.append(contains_key)
-
-            network_kwargs["only_if_contains"] = only_if_contains
-        elif network_config["type"] == "lokr":
-            # find the factor
-            largest_factor = 0
-            for key, value in lora_state_dict.items():
-                if "lokr_w1" in key:
-                    factor = int(value.shape[0])
-                    if factor > largest_factor:
-                        largest_factor = factor
-            network_config["lokr_full_rank"] = True
-            network_config["lokr_factor"] = largest_factor
-
-            only_if_contains = []
-            for key in lora_state_dict.keys():
-                if "lokr_w1" in key:
-                    contains_key = key.split(".lokr_w1")[0]
-                    contains_key = contains_key.replace("lycoris_", "")
-                    if contains_key not in only_if_contains:
-                        only_if_contains.append(contains_key)
-            network_kwargs["only_if_contains"] = only_if_contains
-        
-        if hasattr(base_model, 'target_lora_modules'):
-            network_kwargs['target_lin_modules'] = base_model.target_lora_modules
-
-        # todo auto grab these
-        # get dim and scale
-        network_config = NetworkConfig(**network_config)
-
-        network = LoRASpecialNetwork(
-            text_encoder=None,
-            unet=model_to_quantize,
-            lora_dim=network_config.linear,
-            multiplier=1.0,
-            alpha=network_config.linear_alpha,
-            # conv_lora_dim=self.network_config.conv,
-            # conv_alpha=self.network_config.conv_alpha,
-            train_unet=True,
-            train_text_encoder=False,
-            network_config=network_config,
-            network_type=network_config.type,
-            transformer_only=network_config.transformer_only,
-            is_transformer=base_model.is_transformer,
-            base_model=base_model,
-            is_ara=True,
-            **network_kwargs
-        )
-        network.apply_to(
-            None, model_to_quantize, apply_text_encoder=False, apply_unet=True
-        )
-        network.force_to(base_model.device_torch, dtype=base_model.torch_dtype)
-        network._update_torch_multiplier()
-        network.load_weights(lora_state_dict)
-        network.eval()
-        network.is_active = True
-        network.can_merge_in = False
-        base_model.accuracy_recovery_adapter = network
-
-        # quantize it
-        lora_exclude_modules = []
-        quantization_type = get_qtype(base_model.model_config.qtype)
-        for lora_module in tqdm(network.unet_loras, desc="Attaching quantization"):
-            # the lora has already hijacked the original module
-            orig_module = lora_module.org_module[0]
-            orig_module.to(base_model.torch_dtype)
-            # make the params not require gradients
-            for param in orig_module.parameters():
-                param.requires_grad = False
-            quantize(orig_module, weights=quantization_type)
-            freeze(orig_module)
-            module_name = lora_module.lora_name.replace('$$', '.').replace('transformer.', '')
-            lora_exclude_modules.append(module_name)
-            if base_model.model_config.low_vram:
-                # move it back to cpu
-                orig_module.to("cpu")
-        pass
-        # quantize additional layers
-        print_acc(" - quantizing additional layers")
-        quantization_type = get_qtype('uint8')
-        quantize(
             model_to_quantize,
-            weights=quantization_type,
+            ara_path=base_model.model_config.accuracy_recovery_adapter,
-            exclude=lora_exclude_modules + exclude_modules
+            exclude=exclude_modules,
         )
     else:
+        _quantize_model_blocks(base_model, model_to_quantize, exclude_modules)
+
+
+@torch.no_grad()
+def attach_ara_and_quantize(
+    base_model: "BaseModel",
+    model_to_quantize: torch.nn.Module,
+    ara_path: str,
+    exclude: Optional[List[str]] = None,
+):
+    """Load an accuracy recovery adapter as a live network on the module and
+    quantize around it (adapter-hijacked linears at the configured qtype,
+    everything else uint8). The network lands on
+    base_model.accuracy_recovery_adapter."""
+    from toolkit.config_modules import NetworkConfig
+    from toolkit.lora_special import LoRASpecialNetwork
+
+    exclude_modules = list(exclude or [])
+    load_lora_path = ara_path
+
+    if not os.path.exists(load_lora_path):
+        # not local file, grab from the hub
+
+        path_split = load_lora_path.split("/")
+        if len(path_split) > 3:
+            raise ValueError(
+                "The accuracy recovery adapter path must be a local path or for a hf repo, 'username/repo_name/filename.safetensors'."
+            )
+        repo_id = f"{path_split[0]}/{path_split[1]}"
+        print_acc(f"Grabbing lora from the hub: {load_lora_path}")
+        new_lora_path = hf_hub_download(
+            repo_id,
+            filename=path_split[-1],
+        )
+        # replace the path
+        load_lora_path = new_lora_path
+
+    # build the lora config based on the lora weights
+    lora_state_dict = load_file(load_lora_path)
+        
+    if hasattr(base_model, "convert_lora_weights_before_load"):
+        lora_state_dict = base_model.convert_lora_weights_before_load(lora_state_dict)
+        
+    network_config = {
+        "type": "lora",
+        "network_kwargs": {"only_if_contains": []},
+        "transformer_only": False,
+    }
+    first_key = list(lora_state_dict.keys())[0]
+    first_weight = lora_state_dict[first_key]
+    # if it starts with lycoris and includes lokr
+    if first_key.startswith("lycoris") and any(
+        "lokr" in key for key in lora_state_dict.keys()
+    ):
+        network_config["type"] = "lokr"
+        
+    network_kwargs = {}
+
+    # find firse loraA weight
+    if network_config["type"] == "lora":
+        linear_dim = None
+        for key, value in lora_state_dict.items():
+            if "lora_A" in key:
+                linear_dim = int(value.shape[0])
+                break
+        linear_alpha = linear_dim
+        network_config["linear"] = linear_dim
+        network_config["linear_alpha"] = linear_alpha
+
+        # we build the keys to match every key
+        only_if_contains = []
+        for key in lora_state_dict.keys():
+            contains_key = key.split(".lora_")[0]
+            if contains_key not in only_if_contains:
+                only_if_contains.append(contains_key)
+
+        network_kwargs["only_if_contains"] = only_if_contains
+    elif network_config["type"] == "lokr":
+        # find the factor
+        largest_factor = 0
+        for key, value in lora_state_dict.items():
+            if "lokr_w1" in key:
+                factor = int(value.shape[0])
+                if factor > largest_factor:
+                    largest_factor = factor
+        network_config["lokr_full_rank"] = True
+        network_config["lokr_factor"] = largest_factor
+
+        only_if_contains = []
+        for key in lora_state_dict.keys():
+            if "lokr_w1" in key:
+                contains_key = key.split(".lokr_w1")[0]
+                contains_key = contains_key.replace("lycoris_", "")
+                if contains_key not in only_if_contains:
+                    only_if_contains.append(contains_key)
+        network_kwargs["only_if_contains"] = only_if_contains
+        
+    if hasattr(base_model, 'target_lora_modules'):
+        network_kwargs['target_lin_modules'] = base_model.target_lora_modules
+
+    # todo auto grab these
+    # get dim and scale
+    network_config = NetworkConfig(**network_config)
+
+    network = LoRASpecialNetwork(
+        text_encoder=None,
+        unet=model_to_quantize,
+        lora_dim=network_config.linear,
+        multiplier=1.0,
+        alpha=network_config.linear_alpha,
+        # conv_lora_dim=self.network_config.conv,
+        # conv_alpha=self.network_config.conv_alpha,
+        train_unet=True,
+        train_text_encoder=False,
+        network_config=network_config,
+        network_type=network_config.type,
+        transformer_only=network_config.transformer_only,
+        is_transformer=base_model.is_transformer,
+        base_model=base_model,
+        is_ara=True,
+        **network_kwargs
+    )
+    network.apply_to(
+        None, model_to_quantize, apply_text_encoder=False, apply_unet=True
+    )
+    network.force_to(base_model.device_torch, dtype=base_model.torch_dtype)
+    network._update_torch_multiplier()
+    network.load_weights(lora_state_dict)
+    network.eval()
+    network.is_active = True
+    network.can_merge_in = False
+    base_model.accuracy_recovery_adapter = network
+
+    # quantize it
+    lora_exclude_modules = []
+    quantization_type = get_qtype(base_model.model_config.qtype)
+    for lora_module in tqdm(network.unet_loras, desc="Attaching quantization"):
+        # the lora has already hijacked the original module
+        orig_module = lora_module.org_module[0]
+        orig_module.to(base_model.torch_dtype)
+        # make the params not require gradients
+        for param in orig_module.parameters():
+            param.requires_grad = False
+        quantize(orig_module, weights=quantization_type)
+        freeze(orig_module)
+        module_name = lora_module.lora_name.replace('$$', '.').replace('transformer.', '')
+        lora_exclude_modules.append(module_name)
+        if base_model.model_config.low_vram:
+            # move it back to cpu
+            orig_module.to("cpu")
+    pass
+    # quantize additional layers
+    print_acc(" - quantizing additional layers")
+    quantization_type = get_qtype('uint8')
+    quantize(
+        model_to_quantize,
+        weights=quantization_type,
+        exclude=lora_exclude_modules + exclude_modules
+    )
+
+
+def _quantize_model_blocks(base_model, model_to_quantize, exclude_modules):
+    if True:
         # quantize model the original way without an accuracy recovery adapter
         # move and quantize only certain pieces at a time.
         quantization_type = get_qtype(base_model.model_config.qtype)