SDXL: Use VAE on GPU

bghira · Jul 27, 2023 · 0fe3913 · 0fe3913
1 parent b324b19
commit 0fe3913
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/discord_tron_client/classes/image_manipulation/diffusion.py b/discord_tron_client/classes/image_manipulation/diffusion.py
@@ -367,7 +367,7 @@ def get_sdxl_refiner_pipe(self):
             scheduler_config={"name": "fast"},
             model_id=refiner_model,
         )
-        pipeline.vae = AutoencoderKL.from_pretrained('ptx0/s2', subfolder='vae', torch_dtype=self.torch_dtype, use_safetensors=True, use_auth_token=config.get_huggingface_api_key())
+        pipeline.vae = AutoencoderKL.from_pretrained('ptx0/s2', subfolder='vae', torch_dtype=self.torch_dtype, use_safetensors=True, use_auth_token=config.get_huggingface_api_key()).to(self.device)
         return pipeline
 
     def enforce_zero_terminal_snr(self, betas):