This commit is contained in:
lllyasviel
2024-01-25 08:06:15 -08:00
parent a3b4f9df29
commit 1c19f8f303
2 changed files with 5 additions and 0 deletions

View File

@@ -147,6 +147,9 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
[unet_patcher],
unet_patcher.memory_required([x.shape[0] * 2] + list(x.shape[1:])) + inference_memory)
self.model_wrap.log_sigmas = self.model_wrap.log_sigmas.to(unet_patcher.current_device)
self.model_wrap.sigmas = self.model_wrap.sigmas.to(unet_patcher.current_device)
steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps)
sigmas = self.get_sigmas(p, steps)

View File

@@ -104,6 +104,8 @@ class CompVisSampler(sd_samplers_common.Sampler):
[unet_patcher],
unet_patcher.memory_required([x.shape[0] * 2] + list(x.shape[1:])) + inference_memory)
self.model_wrap.inner_model.alphas_cumprod = self.model_wrap.inner_model.alphas_cumprod.to(unet_patcher.current_device)
steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps)
timesteps = self.get_timesteps(p, steps)