further a guess to what to convert image to depending on number of ch…

…annels
lucidrains · Oct 23, 2023 · 5ff2393 · 5ff2393
1 parent 3b9f0f7
commit 5ff2393
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -873,7 +873,7 @@ def __init__(
         amp = False,
         mixed_precision_type = 'fp16',
         split_batches = True,
-        convert_image_to = 'RGB',
+        convert_image_to = None,
         calculate_fid = True,
         inception_block_idx = 2048,
         max_grad_norm = 1.,
@@ -895,6 +895,11 @@ def __init__(
         self.channels = diffusion_model.channels
         is_ddim_sampling = diffusion_model.is_ddim_sampling
 
+        # default convert_image_to depending on channels
+
+        if not exists(convert_image_to):
+            convert_image_to = {1: 'L', 3: 'RGB', 4: 'RGBA'}.get(self.channels)
+
         # sampling and training hyperparameters
 
         assert has_int_squareroot(num_samples), 'number of samples must have an integer square root'

diff --git a/denoising_diffusion_pytorch/version.py b/denoising_diffusion_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.9.3'
+__version__ = '1.9.4'