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

Exclude sensative layers from quantization on wan models

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-10 08:56:05 -06:00
parent 883d60eb71
commit e7951ad29e
3 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -211,6 +211,12 @@ class Wan2214bModel(Wan21):
if not self.train_high_noise or not self.train_low_noise: if not self.train_high_noise or not self.train_low_noise:
self.target_lora_modules = ["WanTransformer3DModel"] self.target_lora_modules = ["WanTransformer3DModel"]
def get_quantization_exclude_modules(self):
# the timestep/text conditioning embedders and the final projection feed
# every downstream modulation; keep them in full precision when quantizing.
# names are relative to each individual transformer (they quantize separately)
return ["condition_embedder*", "proj_out*"]
@property @property
def max_step_saves_to_keep_multiplier(self): def max_step_saves_to_keep_multiplier(self):
# the cleanup mechanism checks this to see how many saves to keep # the cleanup mechanism checks this to see how many saves to keep

5
extensions_built_in/diffusion_models/wan22/wan22_5b_model.py
View File

@@ -118,6 +118,11 @@ class Wan225bModel(Wan21):
# 16x compression and 2x2 patch size # 16x compression and 2x2 patch size
return 32 return 32
def get_quantization_exclude_modules(self):
# the timestep/text conditioning embedders and the final projection feed
# every downstream modulation; keep them in full precision when quantizing
return ["condition_embedder*", "proj_out*"]
def get_generation_pipeline(self): def get_generation_pipeline(self):
# todo unipc got broken in a diffusers update. Use euler for now. # todo unipc got broken in a diffusers update. Use euler for now.
# scheduler = UniPCMultistepScheduler(**self._wan_generation_scheduler_config) # scheduler = UniPCMultistepScheduler(**self._wan_generation_scheduler_config)

3
extensions_built_in/diffusion_models/wan22/wan22_pipeline.py
View File

@@ -88,6 +88,9 @@ class Wan22Pipeline(WanPipeline):
text_encoder_device = self.text_encoder.device text_encoder_device = self.text_encoder.device
device = self._exec_device device = self._exec_device
if vae_device == torch.device("cpu"):
vae_device = self.reset_device_map
if self._aggressive_offload: if self._aggressive_offload:
print("Unloading vae") print("Unloading vae")
self.vae.to("cpu") self.vae.to("cpu")