Exclude sensative layers from quantization on wan models

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

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)

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")