Enable neural network upscalers for highres. fix

This commit is contained in:
AUTOMATIC
2022-09-20 19:32:26 +03:00
parent 2c9777fcc7
commit 06cd206107
3 changed files with 23 additions and 2 deletions

View File

@@ -450,7 +450,27 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
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")
if opts.upscaler_for_hires_fix is None or opts.upscaler_for_hires_fix == "None":
decoded_samples = torch.nn.functional.interpolate(decoded_samples, size=(self.height, self.width), mode="bilinear")
else:
lowres_samples = torch.clamp((decoded_samples + 1.0) / 2.0, min=0.0, max=1.0)
batch_images = []
for i, x_sample in enumerate(lowres_samples):
x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
x_sample = x_sample.astype(np.uint8)
image = Image.fromarray(x_sample)
upscaler = [x for x in shared.sd_upscalers if x.name == opts.upscaler_for_hires_fix][0]
image = upscaler.upscale(image, self.width, self.height)
image = np.array(image).astype(np.float32) / 255.0
image = np.moveaxis(image, 2, 0)
batch_images.append(image)
decoded_samples = torch.from_numpy(np.array(batch_images))
decoded_samples = decoded_samples.to(shared.device)
decoded_samples = 2. * decoded_samples - 1.
samples = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(decoded_samples))
shared.state.nextjob()