Made dummy latents smaller. Minor code cleanups

This commit is contained in:
random_thoughtss 2022-10-20 09:45:03 -07:00
parent aa7ff2a197
commit 92a17a7a4a
2 changed files with 8 additions and 5 deletions

View File

@ -557,7 +557,8 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
else: else:
# Dummy zero conditioning if we're not using inpainting model. # Dummy zero conditioning if we're not using inpainting model.
# Still takes up a bit of memory, but no encoder call. # Still takes up a bit of memory, but no encoder call.
image_conditioning = torch.zeros(x.shape[0], 5, x.shape[-2], x.shape[-1], dtype=x.dtype, device=x.device) # Pretty sure we can just make this a 1x1 image since its not going to be used besides its batch size.
image_conditioning = torch.zeros(x.shape[0], 5, 1, 1, dtype=x.dtype, device=x.device)
return image_conditioning return image_conditioning
@ -759,7 +760,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
self.image_conditioning = self.image_conditioning.to(shared.device).type(self.sd_model.dtype) self.image_conditioning = self.image_conditioning.to(shared.device).type(self.sd_model.dtype)
else: else:
self.image_conditioning = torch.zeros( self.image_conditioning = torch.zeros(
self.init_latent.shape[0], 5, self.init_latent.shape[-2], self.init_latent.shape[-1], self.init_latent.shape[0], 5, 1, 1,
dtype=self.init_latent.dtype, dtype=self.init_latent.dtype,
device=self.init_latent.device device=self.init_latent.device
) )

View File

@ -138,7 +138,7 @@ class VanillaStableDiffusionSampler:
if self.stop_at is not None and self.step > self.stop_at: if self.stop_at is not None and self.step > self.stop_at:
raise InterruptedException raise InterruptedException
# Have to unwrap the inpainting conditioning here to perform pre-preocessing # Have to unwrap the inpainting conditioning here to perform pre-processing
image_conditioning = None image_conditioning = None
if isinstance(cond, dict): if isinstance(cond, dict):
image_conditioning = cond["c_concat"][0] image_conditioning = cond["c_concat"][0]
@ -165,6 +165,8 @@ class VanillaStableDiffusionSampler:
img_orig = self.sampler.model.q_sample(self.init_latent, ts) img_orig = self.sampler.model.q_sample(self.init_latent, ts)
x_dec = img_orig * self.mask + self.nmask * x_dec x_dec = img_orig * self.mask + self.nmask * x_dec
# Wrap the image conditioning back up since the DDIM code can accept the dict directly.
# Note that they need to be lists because it just concatenates them later.
if image_conditioning is not None: if image_conditioning is not None:
cond = {"c_concat": [image_conditioning], "c_crossattn": [cond]} cond = {"c_concat": [image_conditioning], "c_crossattn": [cond]}
unconditional_conditioning = {"c_concat": [image_conditioning], "c_crossattn": [unconditional_conditioning]} unconditional_conditioning = {"c_concat": [image_conditioning], "c_crossattn": [unconditional_conditioning]}