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

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