diff --git a/extensions_built_in/audio_models/ace_step/ace_step_15_model.py b/extensions_built_in/audio_models/ace_step/ace_step_15_model.py index 430f82b..e91c444 100644 --- a/extensions_built_in/audio_models/ace_step/ace_step_15_model.py +++ b/extensions_built_in/audio_models/ace_step/ace_step_15_model.py @@ -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"] diff --git a/extensions_built_in/captioner/prompts/ideogram4_prompt.py b/extensions_built_in/captioner/prompts/ideogram4_prompt.py index 6cd814b..976b363 100644 --- a/extensions_built_in/captioner/prompts/ideogram4_prompt.py +++ b/extensions_built_in/captioner/prompts/ideogram4_prompt.py @@ -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) diff --git a/extensions_built_in/diffusion_models/anima/anima.py b/extensions_built_in/diffusion_models/anima/anima.py index 0bd6c99..c9db405 100644 --- a/extensions_built_in/diffusion_models/anima/anima.py +++ b/extensions_built_in/diffusion_models/anima/anima.py @@ -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 diff --git a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py index 2cb4fde..24d4d60 100644 --- a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py +++ b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py @@ -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) --- diff --git a/extensions_built_in/diffusion_models/chroma/chroma_model.py b/extensions_built_in/diffusion_models/chroma/chroma_model.py index 6a3aa9a..8d5a11a 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_model.py @@ -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) diff --git a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py index 2fcc390..6aacd39 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py @@ -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) diff --git a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py index 497207a..393cb0a 100644 --- a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py +++ b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py @@ -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) diff --git a/extensions_built_in/diffusion_models/ernie_image/transformer.py b/extensions_built_in/diffusion_models/ernie_image/transformer.py index ca3a2ae..642de6d 100644 --- a/extensions_built_in/diffusion_models/ernie_image/transformer.py +++ b/extensions_built_in/diffusion_models/ernie_image/transformer.py @@ -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__( diff --git a/extensions_built_in/diffusion_models/example_model/example_model.py b/extensions_built_in/diffusion_models/example_model/example_model.py index 51b973e..c13bf75 100644 --- a/extensions_built_in/diffusion_models/example_model/example_model.py +++ b/extensions_built_in/diffusion_models/example_model/example_model.py @@ -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") diff --git a/extensions_built_in/diffusion_models/example_model/src/model.py b/extensions_built_in/diffusion_models/example_model/src/model.py index 30fc348..80db21a 100644 --- a/extensions_built_in/diffusion_models/example_model/src/model.py +++ b/extensions_built_in/diffusion_models/example_model/src/model.py @@ -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) diff --git a/extensions_built_in/diffusion_models/f_light/f_light.py b/extensions_built_in/diffusion_models/f_light/f_light.py index 018b97b..647cd88 100644 --- a/extensions_built_in/diffusion_models/f_light/f_light.py +++ b/extensions_built_in/diffusion_models/f_light/f_light.py @@ -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() diff --git a/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py b/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py index f565abe..505d65e 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py @@ -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 diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index ec4c106..cc023b4 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -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") diff --git a/extensions_built_in/diffusion_models/flux2/src/model.py b/extensions_built_in/diffusion_models/flux2/src/model.py index 5c04503..ee7dd58 100644 --- a/extensions_built_in/diffusion_models/flux2/src/model.py +++ b/extensions_built_in/diffusion_models/flux2/src/model.py @@ -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() diff --git a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py index 535e9fb..a67ba7f 100644 --- a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py +++ b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py @@ -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( diff --git a/extensions_built_in/diffusion_models/hidream/hidream_model.py b/extensions_built_in/diffusion_models/hidream/hidream_model.py index 9a1925d..01a378c 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_model.py @@ -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 diff --git a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py index 1b374ec..41bcf0c 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py @@ -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 diff --git a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py index 308892d..c57844c 100644 --- a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py +++ b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py @@ -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) diff --git a/extensions_built_in/diffusion_models/ideogram4/src/transformer.py b/extensions_built_in/diffusion_models/ideogram4/src/transformer.py index 4abc5c8..fce6c43 100644 --- a/extensions_built_in/diffusion_models/ideogram4/src/transformer.py +++ b/extensions_built_in/diffusion_models/ideogram4/src/transformer.py @@ -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 diff --git a/extensions_built_in/diffusion_models/krea2/krea2.py b/extensions_built_in/diffusion_models/krea2/krea2.py index 90462d9..705744d 100644 --- a/extensions_built_in/diffusion_models/krea2/krea2.py +++ b/extensions_built_in/diffusion_models/krea2/krea2.py @@ -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() diff --git a/extensions_built_in/diffusion_models/krea2/src/mmdit.py b/extensions_built_in/diffusion_models/krea2/src/mmdit.py index 5e7cfcb..f4036be 100644 --- a/extensions_built_in/diffusion_models/krea2/src/mmdit.py +++ b/extensions_built_in/diffusion_models/krea2/src/mmdit.py @@ -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 diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index dc561a3..7e947d0 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -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) diff --git a/extensions_built_in/diffusion_models/mageflow/mageflow.py b/extensions_built_in/diffusion_models/mageflow/mageflow.py index d54d628..53429b5 100644 --- a/extensions_built_in/diffusion_models/mageflow/mageflow.py +++ b/extensions_built_in/diffusion_models/mageflow/mageflow.py @@ -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() diff --git a/extensions_built_in/diffusion_models/mageflow/src/transformer.py b/extensions_built_in/diffusion_models/mageflow/src/transformer.py index 7ce1a76..d2ae0e1 100644 --- a/extensions_built_in/diffusion_models/mageflow/src/transformer.py +++ b/extensions_built_in/diffusion_models/mageflow/src/transformer.py @@ -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 diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index c9ce2ed..0c7fb71 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -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() diff --git a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py index ee80c39..1a39802 100644 --- a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py +++ b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py @@ -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( diff --git a/extensions_built_in/diffusion_models/omnigen2/__init__.py b/extensions_built_in/diffusion_models/omnigen2/__init__.py index 8887767..58b6be8 100644 --- a/extensions_built_in/diffusion_models/omnigen2/__init__.py +++ b/extensions_built_in/diffusion_models/omnigen2/__init__.py @@ -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") diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py index 7358b1b..120d7e1 100644 --- a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py @@ -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) diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py index f1ad25e..4c4be72 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py @@ -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) diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py index efab9be..3e15cc4 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py @@ -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 ( diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py index bb3a9c2..871159c 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py @@ -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 ( diff --git a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py index 97ccd52..fa43a58 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py @@ -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 diff --git a/extensions_built_in/diffusion_models/z_image/z_image.py b/extensions_built_in/diffusion_models/z_image/z_image.py index bc2a30e..7ffc823 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image.py +++ b/extensions_built_in/diffusion_models/z_image/z_image.py @@ -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) diff --git a/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py b/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py index 289ca1a..b642935 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py +++ b/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py @@ -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( diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py index d9af932..fcc2ce3 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py @@ -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) diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py index 03995b1..8a0a5f2 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py @@ -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() diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 8f06dfb..131647b 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -1628,37 +1628,41 @@ class BaseModel: else: self.text_encoder.to(*args, **kwargs) - def prepare_text_encoder(self, text_encoder, dtype=None): - """Standard post-load text-encoder policy: layer offloading, device - placement, then quantize_te. Skips quantization when the checkpoint - loaded pre-quantized.""" - from optimum.quanto import freeze - from toolkit.memory_management import MemoryManager - from toolkit.util.quantize import get_qtype, quantize - - dtype = dtype if dtype is not None else self.torch_dtype - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te and not getattr( - text_encoder, "aitk_is_quantized", False - ): - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - return text_encoder - + def component_load_kwargs(self, role: str = "transformer", dtype=None): + """kwargs for a v2 module's .load()/.aitk_post_load(), derived from + model_config: qtype (with the accuracy recovery adapter recombined), + offload fraction, devices, low_vram placement. Roles: "transformer", + "te", "vae".""" + mc = self.model_config + qtype, offload = None, 0.0 + if role == "transformer": + if mc.quantize: + qtype = mc.qtype + if mc.accuracy_recovery_adapter and "|" not in (qtype or ""): + qtype = f"{qtype}|{mc.accuracy_recovery_adapter}" + if mc.layer_offloading: + offload = mc.layer_offloading_transformer_percent + elif role == "te": + if mc.quantize_te: + qtype = mc.qtype_te + if mc.layer_offloading: + offload = mc.layer_offloading_text_encoder_percent + if dtype is None: + dtype = self.vae_torch_dtype if role == "vae" else self.torch_dtype + device = self.te_device_torch if role == "te" else self.device_torch + if mc.low_vram and role in ("transformer", "te"): + device = "cpu" + elif role == "vae": + device = self.vae_device_torch + return dict( + qtype=qtype, + offload=offload, + dtype=dtype, + device=device, + quantize_device=self.device_torch, + base_model=self, + use_comfy_weights=mc.model_kwargs.get("use_comfy_weights", True), + ) def convert_lora_weights_before_save(self, state_dict): # can be overridden in child classes to convert weights before saving diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 52a509e..ca0d6da 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -408,6 +408,40 @@ a legacy scheduler fallback, and the sampler-name pass-through. Full monolith decomposition (per-arch v2 loading with comfy candidates) remains future work, but every resident component is now poolable by the engine. +## Unified load API (2026-08-28) + +Loading policy moved out of the holders into the mixin. The surface: + +- `ModelClassName.load(name_or_path, qtype=..., offload=..., dtype=..., + device=..., ...)` — sourcing (`load_model`) + `aitk_post_load` in one call. +- `module.aitk_post_load(**kwargs)` — the post-load half alone, for holders + whose checkpoint sourcing is custom (state-dict surgery, combined files). + Handles: qtype (incl. `"qtype|ara_path"` accuracy recovery adapters, with + prequantized-skip via `aitk_is_quantized`), block-streamed quantization + (`quantize_module` in toolkit/util/quantize.py), MemoryManager layer + offloading with per-class `get_offload_ignore_modules()` hooks, and device + placement (low_vram parks on cpu). +- `BaseModel.component_load_kwargs(role)` derives the kwargs from + model_config for roles "transformer" / "te" / "vae" (qtype/qtype_te, ARA + recombination, offload percents, low_vram, te/vae devices). + +All holders converted (transformer + TE quantize/offload/placement blocks +deleted): krea2, z_image, qwen_image, chroma ×2, flux_kontext, f_light, +nucleus, zeta_chroma, z_image_l2p, wan21 (+ subclasses), wan22_14b (dual: +per-transformer load, dual-ARA branch kept), boogu, ernie, ideogram4, +mageflow, prx_pixel, omnigen2, flux2 + klein, hidream + hidream_o1, ltx2 + +ltx2.5 (prequantized ConvRot flagged via aitk_is_quantized), minimax_h3, +anima (conditioner rides the transformer quantize flag at qtype_te), +ace_step_15, example_model (template now teaches the new API). +New `get_offload_ignore_modules` hooks: wan (scale_shift tables — also fixes +wan21 offload which previously offloaded them), ltx2 (all four per-block +tables), gemma3 (embed_tokens), gemma4 (embed_tokens + layer_scalar), +ideogram4 (rotary inv_freq + input/cond projections), ernie (x_embedder), +zeta ZImageDCT (pad tokens). New block-name hooks where only the holder had +them: Flux2, MageFlow, ZImageDCT, Ideogram4Transformer2DModel (+ its exclude +list on MageFlow). `prepare_text_encoder` remains only as the legacy-monolith +path. + ## Testing - [x] `testing/test_model_loading.py`: per-arch load + one small sample through diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index c2e9699..def4cbb 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -122,6 +122,11 @@ class OstrisModelMixin: """fnmatch patterns of sensitive modules to keep in full precision.""" return None + def get_offload_ignore_modules(self) -> Optional[List]: + """Module/parameter objects layer offloading must keep resident (pad + tokens, modulation tables, ...). Default: none.""" + return None + # ---- backend hooks: default to the diffusers ModelMixin API. transformers-lib # models (text encoders) override these three to use PreTrainedModel/AutoConfig. @@ -152,6 +157,131 @@ class OstrisModelMixin: # loading # ------------------------------------------------------------------ + @classmethod + def load( + cls, + name_or_path: str, + qtype: Optional[str] = None, + offload: float = 0.0, + dtype: torch.dtype = torch.bfloat16, + device: Optional[torch.device] = None, + base_model=None, + quantize_device: Optional[torch.device] = None, + exclude_quant_modules: Optional[List[str]] = None, + config_path: Optional[str] = None, + config=None, + subfolder: Optional[str] = None, + use_comfy_weights: bool = True, + **kwargs, + ): + """One-call component load. Resolves/downloads the source (comfy + candidates included), builds the module, quantizes it, attaches layer + offloading, and places it — everything a holder used to hand-roll. + + qtype: quantization type; "type|path.safetensors" also attaches an + accuracy recovery adapter (requires base_model for the network + wiring). Pre-quantized checkpoints skip quantization automatically. + offload: layer-offload fraction (0 disables); the class's + get_offload_ignore_modules hook pins sensitive pieces. + device: final placement. With offload active the memory manager owns + staging, so no blanket device move happens. + base_model: the holder — supplies status prints, quantize_kwargs, and + ARA wiring when given. + """ + model = cls.load_model( + name_or_path, + dtype=dtype, + config_path=config_path, + config=config, + subfolder=subfolder, + use_comfy_weights=use_comfy_weights, + **kwargs, + ) + return model.aitk_post_load( + qtype=qtype, + offload=offload, + dtype=dtype, + device=device, + base_model=base_model, + quantize_device=quantize_device, + exclude_quant_modules=exclude_quant_modules, + ) + + def aitk_post_load( + self, + qtype: Optional[str] = None, + offload: float = 0.0, + dtype: torch.dtype = torch.bfloat16, + device: Optional[torch.device] = None, + base_model=None, + quantize_device: Optional[torch.device] = None, + exclude_quant_modules: Optional[List[str]] = None, + **_load_only_kwargs, # tolerate component_load_kwargs() extras + ): + """The post-load half of `load` — quantize (incl. "qtype|ara"), attach + layer offloading, place on device. Callable directly by holders whose + checkpoint sourcing is custom but want the standard policy.""" + from toolkit.print import print_acc + + status_fn = ( + base_model.print_and_status_update if base_model is not None else print_acc + ) + + ara_path = None + if qtype is not None and "|" in qtype: + qtype, ara_path = qtype.split("|", 1) + + if qtype and not getattr(self, "aitk_is_quantized", False): + from toolkit.util.quantize import ( + attach_ara_and_quantize, + quantize_module, + ) + + exclude = list(exclude_quant_modules or []) + list( + self.get_quantization_exclude_modules() or [] + ) + if ara_path is not None: + if base_model is None: + raise ValueError( + "an accuracy recovery adapter needs base_model= for the " + "network wiring" + ) + status_fn("Quantizing with accuracy recovery adapter") + attach_ara_and_quantize(base_model, self, ara_path, exclude=exclude) + else: + status_fn(f"Quantizing ({qtype})") + quantize_module( + self, + qtype, + device=quantize_device if quantize_device is not None else device, + dtype=dtype, + block_names=self.get_transformer_block_names(), + exclude=exclude, + quantize_kwargs=( + base_model.model_config.quantize_kwargs + if base_model is not None + else None + ), + status_fn=status_fn, + ) + self.aitk_is_quantized = True + self.aitk_qtype = qtype + elif qtype and getattr(self, "aitk_is_quantized", False): + status_fn("Checkpoint is pre-quantized; skipping quantization") + + if offload and offload > 0: + from toolkit.memory_management import MemoryManager + + MemoryManager.attach( + self, + device, + offload_percent=offload, + ignore_modules=list(self.get_offload_ignore_modules() or []), + ) + elif device is not None: + self.to(device) + return self + @classmethod def load_model( cls, diff --git a/toolkit/models/v2/diffusion_models/ltx2.py b/toolkit/models/v2/diffusion_models/ltx2.py index 2657336..9cd4aa3 100644 --- a/toolkit/models/v2/diffusion_models/ltx2.py +++ b/toolkit/models/v2/diffusion_models/ltx2.py @@ -21,6 +21,18 @@ class LTX2VideoTransformer3DModel( def get_transformer_block_names(cls): return ["transformer_blocks"] + def get_offload_ignore_modules(self): + # fp32 scale/shift tables (bare tensors) must stay resident + mods = [] + for block in self.transformer_blocks: + mods += [ + block.scale_shift_table, + block.audio_scale_shift_table, + block.video_a2v_cross_attn_scale_shift_table, + block.audio_a2v_cross_attn_scale_shift_table, + ] + return mods + [self.scale_shift_table, self.audio_scale_shift_table] + class LTX2TextConnectors(DiffusersLTX2TextConnectors, OstrisModelMixin): aitk_subfolder = "connectors" diff --git a/toolkit/models/v2/diffusion_models/wan.py b/toolkit/models/v2/diffusion_models/wan.py index 67f62b6..d71cdaf 100644 --- a/toolkit/models/v2/diffusion_models/wan.py +++ b/toolkit/models/v2/diffusion_models/wan.py @@ -80,6 +80,12 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): def get_transformer_block_names(cls): return ["blocks"] + def get_offload_ignore_modules(self): + # fp32 modulation tables must stay resident on the compute device + return [self.scale_shift_table] + [ + block.scale_shift_table for block in self.blocks + ] + @classmethod def convert_state_dict_on_load(cls, state_dict): # original/comfy wan keys -> diffusers layout via diffusers' own diff --git a/toolkit/models/v2/diffusion_models/z_image.py b/toolkit/models/v2/diffusion_models/z_image.py index 794c3ec..e77990b 100644 --- a/toolkit/models/v2/diffusion_models/z_image.py +++ b/toolkit/models/v2/diffusion_models/z_image.py @@ -26,6 +26,9 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix def get_transformer_block_names(cls): return ["layers"] + def get_offload_ignore_modules(self): + return [self.x_pad_token, self.cap_pad_token] + @classmethod def get_quantization_exclude_modules(cls): # sensitive modules kept in full precision (fnmatch patterns on module diff --git a/toolkit/models/v2/text_encoders/gemma3.py b/toolkit/models/v2/text_encoders/gemma3.py index 8391fb3..8e6bd1f 100644 --- a/toolkit/models/v2/text_encoders/gemma3.py +++ b/toolkit/models/v2/text_encoders/gemma3.py @@ -14,6 +14,9 @@ class Gemma3TextEncoder(Gemma3ForConditionalGeneration, OstrisTransformersMixin) # both layouts seen across transformers versions; missing paths skip return ["model.language_model.layers", "language_model.model.layers"] + def get_offload_ignore_modules(self): + return [self.model.language_model.base_model.embed_tokens] + try: from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel @@ -28,5 +31,12 @@ try: def get_transformer_block_names(cls): return ["layers"] + def get_offload_ignore_modules(self): + # layer_scalar is a bare tensor buffer on each decoder layer; the + # manager never enumerates it, so it must ride along explicitly + return [self.embed_tokens] + [ + layer.layer_scalar for layer in self.layers + ] + except ImportError: Gemma4TextEncoder = None diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py index dc55098..73fde4c 100644 --- a/toolkit/models/wan21/wan21.py +++ b/toolkit/models/wan21/wan21.py @@ -6,7 +6,6 @@ from toolkit.accelerator import unwrap_model from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.dequantize import patch_dequantization_on_save -from toolkit.memory_management.manager import MemoryManager from toolkit.models.base_model import BaseModel from toolkit.prompt_utils import PromptEmbeds from transformers import AutoTokenizer, UMT5EncoderModel @@ -42,7 +41,6 @@ from diffusers.pipelines.wan.pipeline_wan import XLA_AVAILABLE from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback from typing import Any, Callable, Dict, List, Optional, Union from toolkit.models.wan21.wan_lora_convert import convert_to_diffusers, convert_to_original -from toolkit.util.quantize import quantize_model from toolkit.models.v2.text_encoders.umt5 import UMT5TextEncoder from toolkit.models.v2.vae.wan import WanVAE from toolkit.metadata import get_meta_for_safetensors @@ -344,23 +342,11 @@ class Wan21(BaseModel): def load_wan_transformer(self, transformer_path, subfolder=None): self.print_and_status_update("Loading transformer") - dtype = self.torch_dtype - transformer = WanTransformer3DModel.load_model( - transformer_path, dtype=dtype, subfolder=subfolder - ) if self.model_config.split_model_over_gpus: raise ValueError( "Splitting model over gpus is not supported for Wan2.1 models") - if self.model_config.low_vram: - # quantize on the device - transformer.to('cpu', dtype=dtype) - flush() - else: - transformer.to(self.device_torch, dtype=dtype) - flush() - if self.model_config.assistant_lora_path is not None or self.model_config.inference_lora_path is not None: raise ValueError( "Assistant LoRA is not supported for Wan2.1 models currently") @@ -369,23 +355,13 @@ class Wan21(BaseModel): raise ValueError( "Loading LoRA is not supported for Wan2.1 models currently") + # load + quantize + offload + placement, all driven by model_config + transformer = WanTransformer3DModel.load( + transformer_path, + subfolder=subfolder, + **self.component_load_kwargs("transformer"), + ) flush() - - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0: - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to('cpu') return transformer @@ -418,8 +394,9 @@ class Wan21(BaseModel): self.print_and_status_update("Loading UMT5EncoderModel") tokenizer = UMT5TextEncoder.load_tokenizer(te_path) - text_encoder = UMT5TextEncoder.load_model(te_path, dtype=dtype) - self.prepare_text_encoder(text_encoder, dtype=dtype) + text_encoder = UMT5TextEncoder.load( + te_path, **self.component_load_kwargs("te") + ) if self.model_config.low_vram: print("Moving transformer back to GPU") diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index db54f58..1f820a1 100644 --- a/toolkit/util/quantize.py +++ b/toolkit/util/quantize.py @@ -261,6 +261,68 @@ def _has_quantizable_linear(module: torch.nn.Module, weights, exclude=None) -> b return False +@torch.no_grad() +def quantize_module( + module: torch.nn.Module, + qtype: str, + device=None, + dtype: torch.dtype = torch.bfloat16, + block_names: Optional[List[str]] = None, + exclude: Optional[List[str]] = None, + quantize_kwargs: Optional[dict] = None, + status_fn=print_acc, +): + """Module-centric quantization: block-streamed (each repeated block moves + to ``device`` for the math and returns to cpu) with a whole-module pass + for the extras. This is the core the per-model loaders call; holders' + quantize_model wraps it with model_config plumbing.""" + from toolkit.dequantize import patch_dequantization_on_save + + patch_dequantization_on_save(module) + quantization_type = get_qtype(qtype) + exclude = list(exclude or []) + quantize_kwargs = quantize_kwargs or {} + + all_blocks: List[torch.nn.Module] = [] + for name in block_names or []: + # name may be a dotted path for models that nest their blocks + block_list = module + for part in name.split("."): + block_list = getattr(block_list, part, None) + if block_list is None: + break + if block_list is not None: + all_blocks += list(block_list) + if all_blocks: + status_fn(f" - quantizing {len(all_blocks)} blocks") + already_quantized = 0 + for block in tqdm(all_blocks): + if not _has_quantizable_linear(block, quantization_type, exclude): + # pre-quantized checkpoint with a matching qtype: nothing in this + # block would change — skip the device round-trip and the dtype + # cast entirely so the load stays byte-identical + already_quantized += 1 + continue + if device is not None: + block.to(device, dtype=dtype, non_blocking=True) + quantize(block, weights=quantization_type, exclude=exclude, **quantize_kwargs) + freeze(block) + # NOT non_blocking: an async D2H allocates the cpu destination in pinned + # memory, which the caching host allocator keeps forever — that silently + # retained a model-sized chunk of host ram + if device is not None: + block.to("cpu") + if already_quantized: + status_fn( + f" - {already_quantized} blocks already quantized with a matching qtype; left untouched" + ) + + status_fn(" - quantizing extras") + quantize(module, weights=quantization_type, exclude=exclude, **quantize_kwargs) + freeze(module) + return module + + def quantize_model( base_model: "BaseModel", model_to_quantize: torch.nn.Module, @@ -279,153 +341,176 @@ def quantize_model( exclude_modules = base_model.get_quantization_exclude_modules() or [] if base_model.model_config.accuracy_recovery_adapter is not None: - from toolkit.config_modules import NetworkConfig - from toolkit.lora_special import LoRASpecialNetwork - - # we need to load and quantize with an accuracy recovery adapter - # todo handle hf repos - load_lora_path = base_model.model_config.accuracy_recovery_adapter - - if not os.path.exists(load_lora_path): - # not local file, grab from the hub - - path_split = load_lora_path.split("/") - if len(path_split) > 3: - raise ValueError( - "The accuracy recovery adapter path must be a local path or for a hf repo, 'username/repo_name/filename.safetensors'." - ) - repo_id = f"{path_split[0]}/{path_split[1]}" - print_acc(f"Grabbing lora from the hub: {load_lora_path}") - new_lora_path = hf_hub_download( - repo_id, - filename=path_split[-1], - ) - # replace the path - load_lora_path = new_lora_path - - # build the lora config based on the lora weights - lora_state_dict = load_file(load_lora_path) - - if hasattr(base_model, "convert_lora_weights_before_load"): - lora_state_dict = base_model.convert_lora_weights_before_load(lora_state_dict) - - network_config = { - "type": "lora", - "network_kwargs": {"only_if_contains": []}, - "transformer_only": False, - } - first_key = list(lora_state_dict.keys())[0] - first_weight = lora_state_dict[first_key] - # if it starts with lycoris and includes lokr - if first_key.startswith("lycoris") and any( - "lokr" in key for key in lora_state_dict.keys() - ): - network_config["type"] = "lokr" - - network_kwargs = {} - - # find firse loraA weight - if network_config["type"] == "lora": - linear_dim = None - for key, value in lora_state_dict.items(): - if "lora_A" in key: - linear_dim = int(value.shape[0]) - break - linear_alpha = linear_dim - network_config["linear"] = linear_dim - network_config["linear_alpha"] = linear_alpha - - # we build the keys to match every key - only_if_contains = [] - for key in lora_state_dict.keys(): - contains_key = key.split(".lora_")[0] - if contains_key not in only_if_contains: - only_if_contains.append(contains_key) - - network_kwargs["only_if_contains"] = only_if_contains - elif network_config["type"] == "lokr": - # find the factor - largest_factor = 0 - for key, value in lora_state_dict.items(): - if "lokr_w1" in key: - factor = int(value.shape[0]) - if factor > largest_factor: - largest_factor = factor - network_config["lokr_full_rank"] = True - network_config["lokr_factor"] = largest_factor - - only_if_contains = [] - for key in lora_state_dict.keys(): - if "lokr_w1" in key: - contains_key = key.split(".lokr_w1")[0] - contains_key = contains_key.replace("lycoris_", "") - if contains_key not in only_if_contains: - only_if_contains.append(contains_key) - network_kwargs["only_if_contains"] = only_if_contains - - if hasattr(base_model, 'target_lora_modules'): - network_kwargs['target_lin_modules'] = base_model.target_lora_modules - - # todo auto grab these - # get dim and scale - network_config = NetworkConfig(**network_config) - - network = LoRASpecialNetwork( - text_encoder=None, - unet=model_to_quantize, - lora_dim=network_config.linear, - multiplier=1.0, - alpha=network_config.linear_alpha, - # conv_lora_dim=self.network_config.conv, - # conv_alpha=self.network_config.conv_alpha, - train_unet=True, - train_text_encoder=False, - network_config=network_config, - network_type=network_config.type, - transformer_only=network_config.transformer_only, - is_transformer=base_model.is_transformer, - base_model=base_model, - is_ara=True, - **network_kwargs - ) - network.apply_to( - None, model_to_quantize, apply_text_encoder=False, apply_unet=True - ) - network.force_to(base_model.device_torch, dtype=base_model.torch_dtype) - network._update_torch_multiplier() - network.load_weights(lora_state_dict) - network.eval() - network.is_active = True - network.can_merge_in = False - base_model.accuracy_recovery_adapter = network - - # quantize it - lora_exclude_modules = [] - quantization_type = get_qtype(base_model.model_config.qtype) - for lora_module in tqdm(network.unet_loras, desc="Attaching quantization"): - # the lora has already hijacked the original module - orig_module = lora_module.org_module[0] - orig_module.to(base_model.torch_dtype) - # make the params not require gradients - for param in orig_module.parameters(): - param.requires_grad = False - quantize(orig_module, weights=quantization_type) - freeze(orig_module) - module_name = lora_module.lora_name.replace('$$', '.').replace('transformer.', '') - lora_exclude_modules.append(module_name) - if base_model.model_config.low_vram: - # move it back to cpu - orig_module.to("cpu") - pass - # quantize additional layers - print_acc(" - quantizing additional layers") - quantization_type = get_qtype('uint8') - quantize( + attach_ara_and_quantize( + base_model, model_to_quantize, - weights=quantization_type, - exclude=lora_exclude_modules + exclude_modules + ara_path=base_model.model_config.accuracy_recovery_adapter, + exclude=exclude_modules, ) else: + _quantize_model_blocks(base_model, model_to_quantize, exclude_modules) + + +@torch.no_grad() +def attach_ara_and_quantize( + base_model: "BaseModel", + model_to_quantize: torch.nn.Module, + ara_path: str, + exclude: Optional[List[str]] = None, +): + """Load an accuracy recovery adapter as a live network on the module and + quantize around it (adapter-hijacked linears at the configured qtype, + everything else uint8). The network lands on + base_model.accuracy_recovery_adapter.""" + from toolkit.config_modules import NetworkConfig + from toolkit.lora_special import LoRASpecialNetwork + + exclude_modules = list(exclude or []) + load_lora_path = ara_path + + if not os.path.exists(load_lora_path): + # not local file, grab from the hub + + path_split = load_lora_path.split("/") + if len(path_split) > 3: + raise ValueError( + "The accuracy recovery adapter path must be a local path or for a hf repo, 'username/repo_name/filename.safetensors'." + ) + repo_id = f"{path_split[0]}/{path_split[1]}" + print_acc(f"Grabbing lora from the hub: {load_lora_path}") + new_lora_path = hf_hub_download( + repo_id, + filename=path_split[-1], + ) + # replace the path + load_lora_path = new_lora_path + + # build the lora config based on the lora weights + lora_state_dict = load_file(load_lora_path) + + if hasattr(base_model, "convert_lora_weights_before_load"): + lora_state_dict = base_model.convert_lora_weights_before_load(lora_state_dict) + + network_config = { + "type": "lora", + "network_kwargs": {"only_if_contains": []}, + "transformer_only": False, + } + first_key = list(lora_state_dict.keys())[0] + first_weight = lora_state_dict[first_key] + # if it starts with lycoris and includes lokr + if first_key.startswith("lycoris") and any( + "lokr" in key for key in lora_state_dict.keys() + ): + network_config["type"] = "lokr" + + network_kwargs = {} + + # find firse loraA weight + if network_config["type"] == "lora": + linear_dim = None + for key, value in lora_state_dict.items(): + if "lora_A" in key: + linear_dim = int(value.shape[0]) + break + linear_alpha = linear_dim + network_config["linear"] = linear_dim + network_config["linear_alpha"] = linear_alpha + + # we build the keys to match every key + only_if_contains = [] + for key in lora_state_dict.keys(): + contains_key = key.split(".lora_")[0] + if contains_key not in only_if_contains: + only_if_contains.append(contains_key) + + network_kwargs["only_if_contains"] = only_if_contains + elif network_config["type"] == "lokr": + # find the factor + largest_factor = 0 + for key, value in lora_state_dict.items(): + if "lokr_w1" in key: + factor = int(value.shape[0]) + if factor > largest_factor: + largest_factor = factor + network_config["lokr_full_rank"] = True + network_config["lokr_factor"] = largest_factor + + only_if_contains = [] + for key in lora_state_dict.keys(): + if "lokr_w1" in key: + contains_key = key.split(".lokr_w1")[0] + contains_key = contains_key.replace("lycoris_", "") + if contains_key not in only_if_contains: + only_if_contains.append(contains_key) + network_kwargs["only_if_contains"] = only_if_contains + + if hasattr(base_model, 'target_lora_modules'): + network_kwargs['target_lin_modules'] = base_model.target_lora_modules + + # todo auto grab these + # get dim and scale + network_config = NetworkConfig(**network_config) + + network = LoRASpecialNetwork( + text_encoder=None, + unet=model_to_quantize, + lora_dim=network_config.linear, + multiplier=1.0, + alpha=network_config.linear_alpha, + # conv_lora_dim=self.network_config.conv, + # conv_alpha=self.network_config.conv_alpha, + train_unet=True, + train_text_encoder=False, + network_config=network_config, + network_type=network_config.type, + transformer_only=network_config.transformer_only, + is_transformer=base_model.is_transformer, + base_model=base_model, + is_ara=True, + **network_kwargs + ) + network.apply_to( + None, model_to_quantize, apply_text_encoder=False, apply_unet=True + ) + network.force_to(base_model.device_torch, dtype=base_model.torch_dtype) + network._update_torch_multiplier() + network.load_weights(lora_state_dict) + network.eval() + network.is_active = True + network.can_merge_in = False + base_model.accuracy_recovery_adapter = network + + # quantize it + lora_exclude_modules = [] + quantization_type = get_qtype(base_model.model_config.qtype) + for lora_module in tqdm(network.unet_loras, desc="Attaching quantization"): + # the lora has already hijacked the original module + orig_module = lora_module.org_module[0] + orig_module.to(base_model.torch_dtype) + # make the params not require gradients + for param in orig_module.parameters(): + param.requires_grad = False + quantize(orig_module, weights=quantization_type) + freeze(orig_module) + module_name = lora_module.lora_name.replace('$$', '.').replace('transformer.', '') + lora_exclude_modules.append(module_name) + if base_model.model_config.low_vram: + # move it back to cpu + orig_module.to("cpu") + pass + # quantize additional layers + print_acc(" - quantizing additional layers") + quantization_type = get_qtype('uint8') + quantize( + model_to_quantize, + weights=quantization_type, + exclude=lora_exclude_modules + exclude_modules + ) + + +def _quantize_model_blocks(base_model, model_to_quantize, exclude_modules): + if True: # quantize model the original way without an accuracy recovery adapter # move and quantize only certain pieces at a time. quantization_type = get_qtype(base_model.model_config.qtype)