fix remaining issues with transparencies when generating and training DALL-E

Author: lucidrains

generate.py

@@ -135,7 +135,7 @@ def exists(val):
     outputs_dir.mkdir(parents = True, exist_ok = True)
 
     for i, image in tqdm(enumerate(outputs), desc = 'saving images'):
-        save_image(image, outputs_dir / f'{i}.jpg', normalize=True)
+        save_image(image, outputs_dir / f'{i}.png', normalize=True)
         with open(outputs_dir / 'caption.txt', 'w') as f:
             f.write(file_name)
 

train_dalle.py

@@ -318,6 +318,7 @@ def cp_path_to_dir(cp_path, tag):
 
 IMAGE_SIZE = vae.image_size
 CHANNELS = vae.channels
+TRANSPARENT = CHANNELS == 4
 IMAGE_MODE = 'RGBA' if CHANNELS == 4 else 'RGB'
 
 # configure OpenAI VAE for float16s