From c1ab63c2841abea2ff4705cb569c596430dd5dac Mon Sep 17 00:00:00 2001 From: Chris Heald Date: Fri, 16 Sep 2022 16:12:52 -0700 Subject: [PATCH] When running a batch with preview turned on, produce a grid of preview images --- scripts/sd_utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scripts/sd_utils.py b/scripts/sd_utils.py index eb35628..e2ca3cc 100644 --- a/scripts/sd_utils.py +++ b/scripts/sd_utils.py @@ -809,7 +809,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)