github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Updated diffusion feature extractor

Browse Source
This commit is contained in:
Jaret Burkett
2025-06-19 15:36:10 -06:00
parent 4586eb5392
commit 8602470952
4 changed files with 44 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
toolkit/stable_diffusion_model.py
View File

@@ -2538,8 +2538,8 @@ class StableDiffusion:
# Move to vae to device if on cpu
if self.vae.device == 'cpu':
self.vae.to(self.device)
latents = latents.to(device, dtype=dtype)
self.vae.to(self.device_torch)
latents = latents.to(self.device_torch, dtype=self.torch_dtype)
latents = (latents / self.vae.config['scaling_factor']) + self.vae.config['shift_factor']
images = self.vae.decode(latents).sample
images = images.to(device, dtype=dtype)