From a558d5b68fc18edf6d2dfc0e178594c92f952766 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 29 Jul 2025 09:13:47 -0600 Subject: [PATCH] Move transformer back to device on aggresive wan 2.2 pipeline after generation. --- .../diffusion_models/wan22/wan22_pipeline.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/extensions_built_in/diffusion_models/wan22/wan22_pipeline.py b/extensions_built_in/diffusion_models/wan22/wan22_pipeline.py index 442b5c7a..6c7b93b7 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_pipeline.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_pipeline.py @@ -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,)