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

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