Merge pull request #80 from neverix/patch-4

Update `show` to work with RGBA images
ai-forever · Dec 6, 2021 · a9ea478 · a9ea478
2 parents eb2488d + c9469e3
commit a9ea478
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/rudalle/pipelines.py b/rudalle/pipelines.py
@@ -100,6 +100,7 @@ def show(pil_images, nrow=4, size=14, save_dir=None, show=True):
         for i, pil_image in enumerate(pil_images):
             pil_image.save(join(save_dir, f'img_{count+i}.png'))
 
+    pil_images = [pil_image.convert('RGB') for pil_image in pil_images]
     imgs = torchvision.utils.make_grid(utils.pil_list_to_torch_tensors(pil_images), nrow=nrow)
     if not isinstance(imgs, list):
         imgs = [imgs.cpu()]