From 89a3994c30d4ab65d86f99395917ba9e20eeaff7 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 23 Aug 2023 09:12:13 +0000 Subject: [PATCH] wrap final sample processing --- modules/processing.py | 28 ++++++++++++++++++++++------ 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/modules/processing.py b/modules/processing.py index 1260e83b0..8142fd717 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -577,6 +577,25 @@ def process_images(p: StableDiffusionProcessing) -> Processed: return res +def validate_sample(sample): + ok = True + try: + sample = sample.astype(np.uint8) + return sample + except Exception as e: + shared.log.error(f'Failed to validate sample values: {e}') + ok = False + if not ok: + try: + sample = np.nan_to_num(sample, nan=0, posinf=255, neginf=0) + sample = sample.astype(np.uint8) + except Exception as e: + shared.log.error(f'Failed to correct sample values: {e}') + sample = np.zeros_like(sample) + sample = sample.astype(np.uint8) + return sample + + def process_images_inner(p: StableDiffusionProcessing) -> Processed: """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" @@ -728,11 +747,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: for i, x_sample in enumerate(x_samples_ddim): p.batch_index = i - if shared.backend == shared.Backend.ORIGINAL: - x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2) - x_sample = x_sample.astype(np.uint8) - else: - x_sample = (255. * x_sample).astype(np.uint8) + x_sample = 255. * (np.moveaxis(x_sample.cpu().numpy(), 0, 2) if shared.backend == shared.Backend.ORIGINAL else x_sample) + x_sample = validate_sample(x_sample) if p.restore_faces: if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_face_restoration: orig = p.restore_faces @@ -954,7 +970,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): 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) + x_sample = validate_sample(x_sample) image = Image.fromarray(x_sample) save_intermediate(image, i) image = images.resize_image(1, image, target_width, target_height, upscaler_name=self.hr_upscaler)