From 0fe3913f384d12430335f9e593a71a5de9d46250 Mon Sep 17 00:00:00 2001 From: bghira Date: Wed, 26 Jul 2023 20:26:30 -0700 Subject: [PATCH] SDXL: Use VAE on GPU --- discord_tron_client/classes/image_manipulation/diffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/discord_tron_client/classes/image_manipulation/diffusion.py b/discord_tron_client/classes/image_manipulation/diffusion.py index d6a8eda2..a8ee10dc 100644 --- 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):