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

32
extensions_built_in/audio_models/ace_step/ace_step_15_model.py
View File

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

4
extensions_built_in/captioner/prompts/ideogram4_prompt.py
View File

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

68
extensions_built_in/diffusion_models/anima/anima.py
View File

@@ -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:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
self.print_and_status_update("Quantizing Text Conditioner")
quantize(text_conditioner, weights=get_qtype(self.model_config.qtype_te))
freeze(text_conditioner)
flush()
# the text conditioner rides the transformer quantize flag (at qtype_te)
# but takes the text-encoder offload/placement policy
tc_kwargs = self.component_load_kwargs("te")
tc_kwargs["qtype"] = (
self.model_config.qtype_te if self.model_config.quantize else None
)
text_conditioner.aitk_post_load(**tc_kwargs)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
)
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
pipe.text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
MemoryManager.attach(
text_conditioner,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
text_conditioner.to("cpu")
else:
transformer.to(self.device_torch, dtype=dtype)
text_conditioner.to(self.device_torch, dtype=dtype)
pipe.text_encoder.to(self.device_torch, dtype=dtype)
# quantize + offload + placement, all driven by model_config
pipe.text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
pipe.text_encoder.requires_grad_(False)
pipe.text_encoder.eval()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(pipe.text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(pipe.text_encoder)
flush()
if self.model_config.low_vram:
pipe.text_encoder.to("cpu")
flush()
flush()
self.noise_scheduler = pipe.scheduler
self.vae = pipe.vae

30
extensions_built_in/diffusion_models/boogu_image/boogu_image.py
View File

@@ -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:
self.print_and_status_update("Quantizing transformer")
quantize_model(self, transformer)
flush()
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
else:
transformer.to(self.device_torch, dtype=dtype)
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
# --- instruction encoder (Qwen3-VL) + processor ---
@@ -203,18 +195,8 @@ class BooguImageModel(BaseModel):
)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing instruction encoder")
text_encoder.to(self.device_torch)
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
if self.model_config.low_vram:
self.print_and_status_update("Moving instruction encoder to CPU")
text_encoder.to("cpu")
else:
text_encoder.to(self.device_torch)
# quantize + offload + placement, all driven by model_config
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
flush()
# --- VAE (FLUX AutoencoderKL) ---

17
extensions_built_in/diffusion_models/chroma/chroma_model.py
View File

@@ -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:
# 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)
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
self.print_and_status_update("Loading T5")
tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype)
self.prepare_text_encoder(text_encoder_2, dtype=dtype)
text_encoder_2 = T5TextEncoder.load(
extras_path, **self.component_load_kwargs("te")
)
# self.print_and_status_update("Loading CLIP")
text_encoder = FakeCLIP(device=self.device_torch)

17
extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py
View File

@@ -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:
# 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)
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
self.print_and_status_update("Loading T5")
tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype)
self.prepare_text_encoder(text_encoder_2, dtype=dtype)
text_encoder_2 = T5TextEncoder.load(
extras_path, **self.component_load_kwargs("te")
)
# self.print_and_status_update("Loading CLIP")
text_encoder = FakeCLIP(device=self.device_torch)

52
extensions_built_in/diffusion_models/ernie_image/ernie_image.py
View File

@@ -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)
if self.model_config.quantize:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[
transformer.x_embedder,
],
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
# load + quantize + offload + placement, all driven by model_config
transformer = ErnieImageTransformer2DModel.load(
model_path, **self.component_load_kwargs("transformer")
)
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(
base_model_path, subfolder="text_encoder", dtype=dtype
text_encoder = Mistral3ModelEncoder.load(
base_model_path, subfolder="text_encoder", **self.component_load_kwargs("te")
)
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
text_encoder.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
self.print_and_status_update("Loading VAE")
vae = Flux2KLVAE.load_model( base_model_path, dtype=dtype
).to(self.device_torch, dtype=dtype)

3
extensions_built_in/diffusion_models/ernie_image/transformer.py
View File

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

23
extensions_built_in/diffusion_models/example_model/example_model.py
View File

@@ -32,7 +32,7 @@ from toolkit.models.base_model import BaseModel
from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from toolkit.util.quantize import quantize, get_qtype
from .src.model import ExampleTransformer2DModel
from .src.pipeline import ExamplePipeline, pad_prompt_embeds
@@ -165,21 +165,14 @@ class ExampleModel(BaseModel):
del state_dict
flush() # gc + empty cuda cache; call it after dropping anything big
if self.model_config.quantize:
# quantize_model handles qtype selection, exclusions and device
# juggling, and leaves the model on CPU
self.print_and_status_update("Quantizing transformer")
quantize_model(self, transformer)
flush()
if self.model_config.low_vram:
# leave it on CPU; get_noise_prediction moves it over when needed
transformer.to("cpu")
else:
transformer.to(self.device_torch, dtype=dtype)
# quantize + offload + placement, all driven by model_config:
# component_load_kwargs derives qtype (incl. an accuracy recovery
# adapter), the layer-offload fraction and the target device
# (low_vram parks on CPU); aitk_post_load applies them. Models whose
# checkpoint sourcing is standard can collapse the build + this into
# one call: ExampleTransformer2DModel.load(path, **kwargs).
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
# For partial layer offloading support see MemoryManager.attach usage
# in ../ideogram4/ideogram4.py or ../z_image/z_image.py.
# --- text encoder + tokenizer (stock transformers model) ---
self.print_and_status_update("Loading text encoder")

24
extensions_built_in/diffusion_models/example_model/src/model.py
View File

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

22
extensions_built_in/diffusion_models/f_light/f_light.py
View File

@@ -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.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)
transformer = DiT.load(
model_path, **self.component_load_kwargs("transformer")
)
flush()
self.print_and_status_update("Loading T5")
tokenizer = T5TextEncoder.load_tokenizer(extras_path, subfolder="tokenizer")
text_encoder = T5TextEncoder.load_model(
extras_path, dtype=dtype, subfolder="text_encoder"
text_encoder = T5TextEncoder.load(
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()

27
extensions_built_in/diffusion_models/flux2/flux2_klein_model.py
View File

@@ -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(
self.flux2_klein_te_path, dtype=dtype, subfolder=""
# load + quantize + offload + placement, all driven by model_config
text_encoder = Qwen3TextEncoder.load(
self.flux2_klein_te_path, subfolder="", **self.component_load_kwargs("te")
)
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Qwen3")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
elif not self.model_config.low_vram:
text_encoder.to(self.device_torch, dtype=dtype)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
flush()
tokenizer = Qwen2Tokenizer.from_pretrained(self.flux2_klein_te_path)
return text_encoder, tokenizer

54
extensions_built_in/diffusion_models/flux2/flux2_model.py
View File

@@ -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 toolkit.util.quantize import quantize, get_qtype, quantize_model
from optimum.quanto import QTensor
from transformers import AutoProcessor, Mistral3ForConditionalGeneration
from toolkit.models.v2.text_encoders.mistral3 import Mistral3TextEncoder
@@ -100,29 +97,12 @@ class Flux2Model(BaseModel):
dtype = self.torch_dtype
self.print_and_status_update("Loading Mistral")
text_encoder = Mistral3TextEncoder.load_model(
MISTRAL_PATH, dtype=dtype, subfolder=""
# load + quantize + offload + placement, all driven by model_config
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:
# patch the state dict method
patch_dequantization_on_save(transformer)
# Avoid full-model peak VRAM allocation before quantization.
self.print_and_status_update("Keeping transformer on CPU for quantization")
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
else:
transformer.to(self.device_torch, dtype=dtype)
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
text_encoder, tokenizer = self.load_te()
self.print_and_status_update("Loading VAE")

4
extensions_built_in/diffusion_models/flux2/src/model.py
View File

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

23
extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py
View File

@@ -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.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)
transformer = FluxTransformer2DModel.load(
model_path, **self.component_load_kwargs("transformer")
)
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)
self.prepare_text_encoder(text_encoder_2, dtype=dtype)
text_encoder_2 = T5TextEncoder.load(
base_model_path, **self.component_load_kwargs("te")
)
self.print_and_status_update("Loading CLIP")
text_encoder = CLIPTextEncoder.load_model(

73
extensions_built_in/diffusion_models/hidream/hidream_model.py
View File

@@ -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(
model_path, dtype=torch.bfloat16
transformer = self.hidream_transformer_class.load(
model_path, **self.component_load_kwargs("transformer")
)
if not self.low_vram:
transformer.to(self.device_torch, dtype=dtype)
if self.model_config.quantize:
self.print_and_status_update("Quantizing transformer")
quantization_type = get_qtype(self.model_config.qtype)
if self.low_vram:
# move and quantize only certain pieces at a time.
all_blocks = list(transformer.double_stream_blocks) + list(transformer.single_stream_blocks)
self.print_and_status_update(" - quantizing transformer blocks")
for block in tqdm(all_blocks):
block.to(self.device_torch, dtype=dtype)
quantize(block, weights=quantization_type)
freeze(block)
block.to('cpu')
# flush()
self.print_and_status_update(" - quantizing extras")
transformer.to(self.device_torch, dtype=dtype)
quantize(transformer, weights=quantization_type)
freeze(transformer)
else:
quantize(transformer, weights=quantization_type)
freeze(transformer)
if self.low_vram:
# unload it for now
transformer.to('cpu')
flush()
self.print_and_status_update("Loading vae")
@@ -183,16 +143,11 @@ class HidreamModel(BaseModel):
flush()
self.print_and_status_update("Loading T5 encoders")
text_encoder_3 = T5TextEncoder.load_model(
extras_path, dtype=torch.bfloat16, subfolder="text_encoder_3"
).to(self.device_torch, dtype=dtype)
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing T5")
quantization_type = get_qtype(self.model_config.qtype_te)
quantize(text_encoder_3, weights=quantization_type)
freeze(text_encoder_3)
flush()
# load + quantize + offload + placement, all driven by model_config
text_encoder_3 = T5TextEncoder.load(
extras_path, subfolder="text_encoder_3", **self.component_load_kwargs("te")
)
flush()
tokenizer_3 = T5TextEncoder.load_tokenizer(
extras_path, subfolder="tokenizer_3", use_fast=False

23
extensions_built_in/diffusion_models/hidream/hidream_o1_model.py
View File

@@ -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:
transformer.to(self.device_torch)
if self.model_config.quantize:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[],
)
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
# move over to device now if low vram

57
extensions_built_in/diffusion_models/ideogram4/ideogram4.py
View File

@@ -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 toolkit.util.quantize import quantize, get_qtype, quantize_model
from optimum.quanto import freeze, QTensor
from optimum.quanto import QTensor
import huggingface_hub
from huggingface_hub.errors import EntryNotFoundError
@@ -365,58 +363,13 @@ class Ideogram4Model(BaseModel):
transformer = self._load_transformer(base)
if self.model_config.quantize:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
else:
transformer.to(self.device_torch, dtype=dtype)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[
transformer.rotary_emb.inv_freq,
transformer.input_proj,
transformer.llm_cond_proj,
],
)
elif self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
else:
# quantize_model leaves the model on CPU; make sure it lands on device.
transformer.to(self.device_torch)
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
tokenizer, text_encoder = self._load_text_encoder(base)
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
text_encoder.to(self.device_torch)
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
elif self.model_config.low_vram:
self.print_and_status_update("Moving text encoder to CPU")
text_encoder.to("cpu")
else:
self.print_and_status_update("Moving text encoder to device")
text_encoder.to(self.device_torch)
# quantize + offload + placement, all driven by model_config
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
flush()
vae = self._load_vae(base)

11
extensions_built_in/diffusion_models/ideogram4/src/transformer.py
View File

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

50
extensions_built_in/diffusion_models/krea2/krea2.py
View File

@@ -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:
self.print_and_status_update("Quantizing transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[
module
for module in transformer.modules()
if isinstance(module, (SimpleModulation, DoubleSharedModulation))
],
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
else:
transformer.to(self.device_torch, dtype=dtype)
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
tokenizer, processor, vl_processor, text_encoder = self._load_text_encoder()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing text encoder")
text_encoder.to(self.device_torch)
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
if self.model_config.low_vram:
self.print_and_status_update("Moving text encoder to CPU")
text_encoder.to("cpu")
else:
text_encoder.to(self.device_torch)
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
flush()
vae = self._load_vae()

8
extensions_built_in/diffusion_models/krea2/src/mmdit.py
View File

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

121
extensions_built_in/diffusion_models/ltx2/ltx2.py
View File

@@ -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:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
ignore_modules = []
for block in transformer.transformer_blocks:
ignore_modules.append(block.scale_shift_table)
ignore_modules.append(block.audio_scale_shift_table)
ignore_modules.append(block.video_a2v_cross_attn_scale_shift_table)
ignore_modules.append(block.audio_a2v_cross_attn_scale_shift_table)
ignore_modules.append(transformer.scale_shift_table)
ignore_modules.append(transformer.audio_scale_shift_table)
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=ignore_modules,
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
@@ -441,25 +413,8 @@ class LTX2Model(BaseModel):
text_encoder.model.vision_tower = None
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
ignore_modules=[
text_encoder.model.language_model.base_model.embed_tokens
],
)
# quantize + offload + placement, all driven by model_config
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
text_encoder.to(self.device_torch, dtype=dtype)
flush()
@@ -1375,38 +1330,11 @@ class LTX25Model(LTX2Model):
pin_stored_fp32(transformer)
flush()
if self.model_config.quantize:
if num_quantized_dit:
self.print_and_status_update(
"Transformer is pre-quantized (ConvRot); skipping quantize"
)
else:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
ignore_modules = []
for block in transformer.transformer_blocks:
ignore_modules.append(block.scale_shift_table)
ignore_modules.append(block.audio_scale_shift_table)
ignore_modules.append(block.video_a2v_cross_attn_scale_shift_table)
ignore_modules.append(block.audio_a2v_cross_attn_scale_shift_table)
ignore_modules.append(transformer.scale_shift_table)
ignore_modules.append(transformer.audio_scale_shift_table)
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=ignore_modules,
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
if num_quantized_dit:
# pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize
transformer.aitk_is_quantized = True
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
self.print_and_status_update("Loading connectors")
@@ -1433,32 +1361,11 @@ class LTX25Model(LTX2Model):
del te_state_dict
flush()
if self.model_config.quantize_te:
if num_quantized_te:
self.print_and_status_update(
"Text encoder is pre-quantized (ConvRot); skipping quantize"
)
else:
self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
# layer_scalar is a bare tensor buffer on each decoder layer; the
# manager never enumerates it, so it must ride along explicitly
ignore_modules = [text_encoder.embed_tokens]
for layer in text_encoder.layers:
ignore_modules.append(layer.layer_scalar)
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
ignore_modules=ignore_modules,
)
if num_quantized_te:
# pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize
text_encoder.aitk_is_quantized = True
# quantize + offload + placement, all driven by model_config
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
text_encoder.to(self.device_torch)
text_encoder.requires_grad_(False)

48
extensions_built_in/diffusion_models/mageflow/mageflow.py
View File

@@ -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:
self.print_and_status_update("Quantizing transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
else:
transformer.to(self.device_torch, dtype=dtype)
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
tokenizer, vl_processor, text_encoder = self._load_text_encoder()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing text encoder")
text_encoder.to(self.device_torch)
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
if self.model_config.low_vram:
self.print_and_status_update("Moving text encoder to CPU")
text_encoder.to("cpu")
else:
text_encoder.to(self.device_torch)
# quantize + offload + placement, all driven by model_config
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
flush()
vae = self._load_vae()

20
extensions_built_in/diffusion_models/mageflow/src/transformer.py
View File

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

55
extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
View File

@@ -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:
self.print_and_status_update("Quantizing transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
)
if self.model_config.low_vram:
self.print_and_status_update("Keeping transformer on CPU")
transformer.to("cpu")
else:
transformer.to(self.device_torch)
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
tokenizer, processor, text_encoder = self._load_text_encoder()
te_prequantized = any(
isinstance(m, OstrisLinear) for m in text_encoder.modules()
)
if self.model_config.quantize_te and not te_prequantized:
self.print_and_status_update("Quantizing text encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
elif self.model_config.quantize_te:
self.print_and_status_update(
"Text encoder is already nvfp4/int8 quantized; skipping quantize_te"
)
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
if self.model_config.low_vram:
text_encoder.to("cpu")
else:
text_encoder.to(self.device_torch)
if any(isinstance(m, OstrisLinear) for m in text_encoder.modules()):
# already nvfp4/int8 quantized; aitk_post_load skips quantize_te
text_encoder.aitk_is_quantized = True
# quantize + offload + placement, all driven by model_config
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
flush()
vae_bundle = self._load_vaes()

29
extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py
View File

@@ -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:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[
],
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
@@ -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)
self.prepare_text_encoder(text_encoder, dtype=dtype)
text_encoder = Qwen3VLTextEncoder.load(
base_model_path, **self.component_load_kwargs("te")
)
self.print_and_status_update("Loading VAE")
vae = QwenImageVAE.load_model(

34
extensions_built_in/diffusion_models/omnigen2/__init__.py
View File

@@ -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(
extras_path, dtype=torch.bfloat16, subfolder="mllm"
# load + quantize + offload + placement, all driven by model_config
mllm = Qwen25VLTextEncoder.load(
extras_path, subfolder="mllm", **self.component_load_kwargs("te")
)
mllm.to(self.device_torch, dtype=dtype)
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Qwen2.5 VL model")
quantization_type = get_qtype(self.model_config.qtype_te)
quantize(mllm, weights=quantization_type)
freeze(mllm)
if self.low_vram:
# unload it for now
mllm.to("cpu")
flush()
self.print_and_status_update("Loading transformer")
transformer = OmniGen2Transformer2DModel.load_model(
model_path, dtype=torch.bfloat16
transformer = OmniGen2Transformer2DModel.load(
model_path, **self.component_load_kwargs("transformer")
)
if not self.low_vram:
transformer.to(self.device_torch, dtype=dtype)
if self.model_config.quantize:
self.print_and_status_update("Quantizing transformer")
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")

32
extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py
View File

@@ -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)
transformer.to(dtype=dtype)
flush()
if self.model_config.quantize:
self.print_and_status_update("Quantizing transformer")
quantize_model(self, transformer)
flush()
if self.model_config.low_vram:
transformer.to("cpu")
else:
transformer.to(self.device_torch, dtype=dtype)
# load + quantize + offload + placement, all driven by model_config
transformer = PRXTransformer2DModel.load(
model_path, **self.component_load_kwargs("transformer")
)
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.to(self.te_device_torch)
text_encoder = Qwen3VLTextOnlyEncoder.load(
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)

35
extensions_built_in/diffusion_models/qwen_image/qwen_image.py
View File

@@ -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(
model_path,
dtype=model_dtype,
use_comfy_weights=self.model_config.model_kwargs.get(
"use_comfy_weights", True
),
# load + quantize + offload + placement, all driven by model_config
transformer = QwenImageTransformer2DModel.load(
model_path, **self.component_load_kwargs("transformer")
)
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)

1
extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py
View File

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

1
extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py
View File

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

78
extensions_built_in/diffusion_models/wan22/wan22_14b_model.py
View File

@@ -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_model(
transformer_path_1, dtype=dtype, subfolder=subfolder_1
transformer_1 = WanTransformer3DModel.load(
transformer_path_1, subfolder=subfolder_1, **per_kwargs
)
flush()
if self.model_config.low_vram:
# quantize on the device
transformer_1.to('cpu', dtype=dtype)
flush()
else:
transformer_1.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize and self.model_config.accuracy_recovery_adapter is None:
# todo handle two ARAs
self.print_and_status_update("Quantizing Transformer 1")
quantize_model(self, transformer_1)
flush()
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer 1 to CPU")
transformer_1.to("cpu")
else:
transformer_1.to(self.device_torch)
self.print_and_status_update("Loading transformer 2")
dtype = self.torch_dtype
transformer_2 = WanTransformer3DModel.load_model(
transformer_path_2, dtype=dtype, subfolder=subfolder_2
transformer_2 = WanTransformer3DModel.load(
transformer_path_2, subfolder=subfolder_2, **per_kwargs
)
flush()
if self.model_config.low_vram:
# quantize on the device
transformer_2.to('cpu', dtype=dtype)
flush()
else:
transformer_2.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize and self.model_config.accuracy_recovery_adapter is None:
# todo handle two ARAs
self.print_and_status_update("Quantizing Transformer 2")
quantize_model(self, transformer_2)
flush()
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer 2 to CPU")
transformer_2.to("cpu")
else:
transformer_2.to(self.device_torch)
layer_offloading_transformer = self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0
# make the combined model
self.print_and_status_update("Creating DualWanTransformer3DModel")
@@ -362,18 +325,13 @@ class Wan2214bModel(Wan21):
if layer_offloading_transformer:
MemoryManager.attach(
transformer_1,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[transformer_1.scale_shift_table] + [block.scale_shift_table for block in transformer_1.blocks]
)
MemoryManager.attach(
transformer_2,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[transformer_2.scale_shift_table] + [block.scale_shift_table for block in transformer_2.blocks]
)
for t in (transformer_1, transformer_2):
MemoryManager.attach(
t,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=t.get_offload_ignore_modules(),
)
return transformer

43
extensions_built_in/diffusion_models/z_image/z_image.py
View File

@@ -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
# into full precision weights first (or the checkpoint was pre-quantized)
if self.model_config.quantize and not transformer.aitk_is_quantized:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[
transformer.x_pad_token,
transformer.cap_pad_token,
],
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
self.print_and_status_update("Text Encoder")
tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
self.prepare_text_encoder(text_encoder, dtype=dtype)
text_encoder = Qwen3TextEncoder.load(
base_model_path, **self.component_load_kwargs("te")
)
self.print_and_status_update("Loading VAE")
vae = KLVAE.load_model(base_model_path, dtype=dtype)

31
extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py
View File

@@ -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:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[
transformer.x_pad_token,
transformer.cap_pad_token,
],
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
self.print_and_status_update("Text Encoder")
tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
self.prepare_text_encoder(text_encoder, dtype=dtype)
text_encoder = Qwen3TextEncoder.load(
base_model_path, **self.component_load_kwargs("te")
)
self.print_and_status_update("Loading VAE")
# vae = AutoencoderKL.from_pretrained(

33
extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py
View File

@@ -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)
if self.model_config.quantize:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[
transformer.x_pad_token,
transformer.cap_pad_token,
],
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
flush()
self.print_and_status_update("Text Encoder")
tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
self.prepare_text_encoder(text_encoder, dtype=dtype)
text_encoder = Qwen3TextEncoder.load(
base_model_path, **self.component_load_kwargs("te")
)
self.print_and_status_update("Loading VAE")
vae = FakeVAE(scaling_factor=1.0)

7
extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py
View File

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

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

373
toolkit/util/quantize.py
View File

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