ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-02-20 04:13:57 +00:00
Code Issues 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

@@ -256,6 +256,14 @@ class Wan22Pipeline(WanPipeline):
# Offload all models
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:
return (video,)