Move the offload and quantization to the base model

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

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)
-        self.prepare_text_encoder(text_encoder, dtype=dtype)
+        text_encoder = UMT5TextEncoder.load(
+            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
-        from toolkit.lora_special import LoRASpecialNetwork
-
-        # 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(
+        attach_ara_and_quantize(
+            base_model,
             model_to_quantize,
-            weights=quantization_type,
-            exclude=lora_exclude_modules + exclude_modules
+            ara_path=base_model.model_config.accuracy_recovery_adapter,
+            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)