support vace quant

2026-03-18 22:08:13 +00:00 · 2025-04-16 10:29:21 +08:00
parent e3ac438f5a
commit 6b24748c80
2 changed files with 19 additions and 1 deletions
--- a/diffsynth/pipelines/wan_video.py
+++ b/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()