RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fixed issue with wan 2.2 getting stuck on CPU

This commit is contained in:
Jaret Burkett 2025-10-09 17:24:25 +00:00
parent 55b8b0e23e
commit 8068755b0a
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
extensions_built_in/diffusion_models/wan22/wan22_14b_model.py
View File

@ -302,6 +302,8 @@ class Wan2214bModel(Wan21):
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer 1 to CPU")
transformer_1.to("cpu")
else:
transformer_1.to(self.device_torch)
self.print_and_status_update("Loading transformer 2")
dtype = self.torch_dtype
@ -327,6 +329,8 @@ class Wan2214bModel(Wan21):
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer 2 to CPU")
transformer_2.to("cpu")
else:
transformer_2.to(self.device_torch)
# make the combined model
self.print_and_status_update("Creating DualWanTransformer3DModel")