github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Move the offload and quantization to the base model

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-28 08:37:56 -06:00
parent 702254688d
commit 85a6880643
45 changed files with 725 additions and 1140 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
toolkit/models/base_model.py
View File

@@ -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

34
toolkit/models/v2/PLANNING.md
View File

@@ -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

130
toolkit/models/v2/_mixin.py
View File

@@ -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,

12
toolkit/models/v2/diffusion_models/ltx2.py
View File

@@ -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"

6
toolkit/models/v2/diffusion_models/wan.py
View File

@@ -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

3
toolkit/models/v2/diffusion_models/z_image.py
View File

@@ -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

10
toolkit/models/v2/text_encoders/gemma3.py
View File

@@ -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

41
toolkit/models/wan21/wan21.py
View File

@@ -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")