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

Move the offload and quantization to the base model

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

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

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