1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00

[sdxl k-diffusion pipeline]move sigma to device (#6757)

move sigma to device

Co-authored-by: yiyixuxu <yixu310@gmail,com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
This commit is contained in:
YiYi Xu
2024-01-31 09:29:15 -10:00
committed by GitHub
parent 2e8d18e699
commit 66722dbea7

View File

@@ -888,10 +888,9 @@ class StableDiffusionXLKDiffusionPipeline(
sigma_min: float = self.k_diffusion_model.sigmas[0].item()
sigma_max: float = self.k_diffusion_model.sigmas[-1].item()
sigmas = get_sigmas_karras(n=num_inference_steps, sigma_min=sigma_min, sigma_max=sigma_max)
sigmas = sigmas.to(device)
else:
sigmas = self.scheduler.sigmas
sigmas = sigmas.to(prompt_embeds.dtype)
sigmas = sigmas.to(dtype=prompt_embeds.dtype, device=device)
# 6. Prepare latent variables
num_channels_latents = self.unet.config.in_channels