Merge pull request #536 from modelscope/wan-vace-quant

support vace quant

diffsynth/pipelines/wan_video.py

@@ -142,6 +142,24 @@ class WanVideoPipeline(BasePipeline):
                     computation_device=self.device,
                 ),
             )
+        if self.vace is not None:
+            enable_vram_management(
+                self.vace,
+                module_map = {
+                    torch.nn.Linear: AutoWrappedLinear,
+                    torch.nn.Conv3d: AutoWrappedModule,
+                    torch.nn.LayerNorm: AutoWrappedModule,
+                    RMSNorm: AutoWrappedModule,
+                },
+                module_config = dict(
+                    offload_dtype=dtype,
+                    offload_device="cpu",
+                    onload_dtype=dtype,
+                    onload_device=self.device,
+                    computation_dtype=self.torch_dtype,
+                    computation_device=self.device,
+                ),
+            )
         self.enable_cpu_offload()
 
 

examples/wanvideo/wan_1.3b_vace.py

@@ -8,7 +8,7 @@ from PIL import Image
 snapshot_download("iic/VACE-Wan2.1-1.3B-Preview", local_dir="models/iic/VACE-Wan2.1-1.3B-Preview")
 
 # Load models
-model_manager = ModelManager(device="cuda")
+model_manager = ModelManager(device="cpu")
 model_manager.load_models(
     [
         "models/iic/VACE-Wan2.1-1.3B-Preview/diffusion_pytorch_model.safetensors",