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

Move the offload and quantization to the base model

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

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

@@ -3,13 +3,11 @@ from typing import List, Optional
import torch
import yaml
from optimum.quanto import freeze
from safetensors.torch import load_file, save_file
from toolkit.accelerator import unwrap_model
from toolkit.basic import flush
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.memory_management import MemoryManager
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.diffusion_models.cosmos import CosmosTransformer3DModel
from toolkit.models.v2.text_encoders.anima import AnimaTextConditioner
@@ -17,7 +15,6 @@ from toolkit.models.v2.text_encoders.qwen3 import Qwen3ModelEncoder
from toolkit.models.v2.vae.qwen_image import QwenImageVAE
from toolkit.prompt_utils import PromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
from toolkit.util.quantize import get_qtype, quantize, quantize_model
try:
from diffusers import AnimaAutoBlocks, AnimaModularPipeline
@@ -287,62 +284,23 @@ class AnimaModel(BaseModel):
transformer = pipe.transformer
text_conditioner = pipe.text_conditioner
if self.model_config.quantize:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
# quantize + offload + placement, all driven by model_config
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
self.print_and_status_update("Quantizing Text Conditioner")
quantize(text_conditioner, weights=get_qtype(self.model_config.qtype_te))
freeze(text_conditioner)
flush()
# the text conditioner rides the transformer quantize flag (at qtype_te)
# but takes the text-encoder offload/placement policy
tc_kwargs = self.component_load_kwargs("te")
tc_kwargs["qtype"] = (
self.model_config.qtype_te if self.model_config.quantize else None
)
text_conditioner.aitk_post_load(**tc_kwargs)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
)
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
pipe.text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
MemoryManager.attach(
text_conditioner,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
text_conditioner.to("cpu")
else:
transformer.to(self.device_torch, dtype=dtype)
text_conditioner.to(self.device_torch, dtype=dtype)
pipe.text_encoder.to(self.device_torch, dtype=dtype)
# quantize + offload + placement, all driven by model_config
pipe.text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
pipe.text_encoder.requires_grad_(False)
pipe.text_encoder.eval()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(pipe.text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(pipe.text_encoder)
flush()
if self.model_config.low_vram:
pipe.text_encoder.to("cpu")
flush()
flush()
self.noise_scheduler = pipe.scheduler
self.vae = pipe.vae