Fix VRAM management (again)

2025-08-03 22:35:24 +03:00
parent 1654741cdf
commit 88b5feb1ed
1 changed files with 4 additions and 0 deletions
@@ -1978,6 +1978,10 @@ class WanVideoSampler:
                    module.offload()
                if hasattr(module, "onload"):
                    module.onload()
+            for block in transformer.blocks:
+                block.modulation = torch.nn.Parameter(block.modulation.to(device))
+            transformer.head.modulation = torch.nn.Parameter(transformer.head.modulation.to(device))
+
        elif model["manual_offloading"]:
            transformer.to(device)