github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Move transformer back to device on aggresive wan 2.2 pipeline after generation.

Browse Source
This commit is contained in:
Jaret Burkett
2025-07-29 09:13:47 -06:00
parent 1d1199b15b
commit a558d5b68f
1 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
extensions_built_in/diffusion_models/wan22/wan22_pipeline.py
View File

@@ -257,6 +257,14 @@ class Wan22Pipeline(WanPipeline):
# Offload all models # Offload all models
self.maybe_free_model_hooks() self.maybe_free_model_hooks()
# move transformer back to device
if self._aggressive_offload:
print("Moving transformer back to device")
self.transformer.to(self._execution_device)
if self.transformer_2 is not None:
self.transformer_2.to(self._execution_device)
flush()
if not return_dict: if not return_dict:
return (video,) return (video,)