Merge branch 'dev' into fix-preview-error

colemickens · Sep 16, 2022 · 419f3f9 · 419f3f9
2 parents 0cf3d1b + ac99b65
commit 419f3f9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/scripts/img2img.py b/scripts/img2img.py
@@ -291,7 +291,7 @@ def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name)
 					correction_target,
 					channel_axis=2
 				    ), cv2.COLOR_LAB2RGB).astype("uint8"))
-				if mask_restore is True:
+				if mask_restore is True and init_mask is not None:
 					color_mask = init_mask.filter(ImageFilter.GaussianBlur(mask_blur_strength))
 					color_mask = color_mask.convert('L')
 					source_image = input_image.convert('RGB')

diff --git a/scripts/sd_utils.py b/scripts/sd_utils.py
@@ -811,7 +811,11 @@ def generation_callback(img, i=0):
 
         x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)  
 
-        pil_image = transforms.ToPILImage()(x_samples_ddim.squeeze_(0)) 			
+        if x_samples_ddim.ndimension() == 4:
+            pil_images = [transforms.ToPILImage()(x.squeeze_(0)) for x in x_samples_ddim]
+            pil_image = image_grid(pil_images, 1)
+        else:
+            pil_image = transforms.ToPILImage()(x_samples_ddim.squeeze_(0))
 
         # update image on the UI so we can see the progress
         st.session_state["preview_image"].image(pil_image)