From 31e057d9a379e41b0d711b6663af1c3e8c16a489 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 13 Mar 2025 10:30:27 -0600 Subject: [PATCH] Fixed issue with device placement in some scenereos when doing low vram on wan --- toolkit/models/wan21/wan21.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py index 105191c3..9ef6891f 100644 --- a/toolkit/models/wan21/wan21.py +++ b/toolkit/models/wan21/wan21.py @@ -116,9 +116,11 @@ class AggressiveWanUnloadPipeline(WanPipeline): vae_device = self.vae.device transformer_device = self.transformer.device text_encoder_device = self.text_encoder.device + device = self.transformer.device + print("Unloading vae") self.vae.to("cpu") - self.text_encoder.to(self._execution_device) + self.text_encoder.to(device) # 1. Check inputs. Raise error if not correct self.check_inputs( @@ -136,8 +138,6 @@ class AggressiveWanUnloadPipeline(WanPipeline): self._current_timestep = None self._interrupt = False - device = self._execution_device - # 2. Define call parameters if prompt is not None and isinstance(prompt, str): batch_size = 1