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 aef30b7..be72fb5 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py @@ -211,6 +211,12 @@ class Wan2214bModel(Wan21): if not self.train_high_noise or not self.train_low_noise: 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 def max_step_saves_to_keep_multiplier(self): # the cleanup mechanism checks this to see how many saves to keep diff --git a/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py index 6850166..f231d74 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_5b_model.py @@ -118,6 +118,11 @@ class Wan225bModel(Wan21): # 16x compression and 2x2 patch size 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): # todo unipc got broken in a diffusers update. Use euler for now. # scheduler = UniPCMultistepScheduler(**self._wan_generation_scheduler_config) diff --git a/extensions_built_in/diffusion_models/wan22/wan22_pipeline.py b/extensions_built_in/diffusion_models/wan22/wan22_pipeline.py index d119343..1d7b797 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_pipeline.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_pipeline.py @@ -88,6 +88,9 @@ class Wan22Pipeline(WanPipeline): text_encoder_device = self.text_encoder.device device = self._exec_device + if vae_device == torch.device("cpu"): + vae_device = self.reset_device_map + if self._aggressive_offload: print("Unloading vae") self.vae.to("cpu")