Skip to content

Commit

Permalink
[sdxl k-diffusion pipeline]move sigma to device (#6757)
Browse files Browse the repository at this point in the history
move sigma to device

Co-authored-by: yiyixuxu <yixu310@gmail,com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
  • Loading branch information
3 people authored Jan 31, 2024
1 parent 2e8d18e commit 66722db
Showing 1 changed file with 1 addition and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -888,10 +888,9 @@ def __call__(
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
Expand Down

0 comments on commit 66722db

Please sign in to comment.