From 6d7ca54a1a9f448419acb31a54c5e28f3e4bcc4c Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Mon, 19 Sep 2022 16:42:56 +0300 Subject: added highres fix feature --- modules/processing.py | 80 ++++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 66 insertions(+), 14 deletions(-) (limited to 'modules/processing.py') diff --git a/modules/processing.py b/modules/processing.py index a0f0e575..e146524d 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -74,11 +74,12 @@ class StableDiffusionProcessing: self.overlay_images = overlay_images self.paste_to = None self.color_corrections = None + self.denoising_strength: float = 0 - def init(self, seed): + def init(self, all_prompts, all_seeds, all_subseeds): pass - def sample(self, x, conditioning, unconditional_conditioning): + def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength): raise NotImplementedError() @@ -303,7 +304,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed: precision_scope = torch.autocast if cmd_opts.precision == "autocast" else contextlib.nullcontext ema_scope = (contextlib.nullcontext if cmd_opts.lowvram else p.sd_model.ema_scope) with torch.no_grad(), precision_scope("cuda"), ema_scope(): - p.init(seed=all_seeds[0]) + p.init(all_prompts, all_seeds, all_subseeds) if state.job_count == -1: state.job_count = p.n_iter @@ -328,13 +329,10 @@ def process_images(p: StableDiffusionProcessing) -> Processed: for comment in model_hijack.comments: comments[comment] = 1 - # we manually generate all input noises because each one should have a specific seed - x = create_random_tensors([opt_C, p.height // opt_f, p.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w, p=p) - if p.n_iter > 1: shared.state.job = f"Batch {n+1} out of {p.n_iter}" - samples_ddim = p.sample(x=x, conditioning=c, unconditional_conditioning=uc) + samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength) if state.interrupted: # if we are interruped, sample returns just noise @@ -406,13 +404,64 @@ def process_images(p: StableDiffusionProcessing) -> Processed: class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): sampler = None + firstphase_width = 0 + firstphase_height = 0 + firstphase_width_truncated = 0 + firstphase_height_truncated = 0 + + def __init__(self, enable_hr=False, scale_latent=True, denoising_strength=0.75, **kwargs): + super().__init__(**kwargs) + self.enable_hr = enable_hr + self.scale_latent = scale_latent + self.denoising_strength = denoising_strength + + def init(self, all_prompts, all_seeds, all_subseeds): + if self.enable_hr: + if state.job_count == -1: + state.job_count = self.n_iter * 2 + else: + state.job_count = state.job_count * 2 + + desired_pixel_count = 512 * 512 + actual_pixel_count = self.width * self.height + scale = math.sqrt(desired_pixel_count / actual_pixel_count) + + self.firstphase_width = math.ceil(scale * self.width / 64) * 64 + self.firstphase_height = math.ceil(scale * self.height / 64) * 64 + self.firstphase_width_truncated = int(scale * self.width) + self.firstphase_height_truncated = int(scale * self.height) + + def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength): + self.sampler = samplers[self.sampler_index].constructor(self.sd_model) + + if not self.enable_hr: + x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self) + samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning) + return samples + + x = create_random_tensors([opt_C, self.firstphase_height // opt_f, self.firstphase_width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self) + samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning) + + truncate_x = (self.firstphase_width - self.firstphase_width_truncated) // opt_f + truncate_y = (self.firstphase_height - self.firstphase_height_truncated) // opt_f + + samples = samples[:, :, truncate_y//2:samples.shape[2]-truncate_y//2, truncate_x//2:samples.shape[3]-truncate_x//2] + + if self.scale_latent: + samples = torch.nn.functional.interpolate(samples, size=(self.height // opt_f, self.width // opt_f), mode="bilinear") + else: + decoded_samples = self.sd_model.decode_first_stage(samples) + decoded_samples = torch.nn.functional.interpolate(decoded_samples, size=(self.height, self.width), mode="bilinear") + samples = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(decoded_samples)) + + shared.state.nextjob() - def init(self, seed): self.sampler = samplers[self.sampler_index].constructor(self.sd_model) + noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self) + samples = self.sampler.sample_img2img(self, samples, noise, conditioning, unconditional_conditioning, steps=self.steps) + + return samples - def sample(self, x, conditioning, unconditional_conditioning): - samples_ddim = self.sampler.sample(self, x, conditioning, unconditional_conditioning) - return samples_ddim class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): sampler = None @@ -435,7 +484,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): self.mask = None self.nmask = None - def init(self, seed): + def init(self, all_prompts, all_seeds, all_subseeds): self.sampler = samplers_for_img2img[self.sampler_index].constructor(self.sd_model) crop_region = None @@ -529,12 +578,15 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): self.mask = torch.asarray(1.0 - latmask).to(shared.device).type(self.sd_model.dtype) self.nmask = torch.asarray(latmask).to(shared.device).type(self.sd_model.dtype) + # this needs to be fixed to be done in sample() using actual seeds for batches if self.inpainting_fill == 2: - self.init_latent = self.init_latent * self.mask + create_random_tensors(self.init_latent.shape[1:], [seed + x + 1 for x in range(self.init_latent.shape[0])]) * self.nmask + self.init_latent = self.init_latent * self.mask + create_random_tensors(self.init_latent.shape[1:], all_seeds[0:self.init_latent.shape[0]]) * self.nmask elif self.inpainting_fill == 3: self.init_latent = self.init_latent * self.mask - def sample(self, x, conditioning, unconditional_conditioning): + def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength): + x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self) + samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning) if self.mask is not None: -- cgit v1.2.1