unet_number on decoder trainer only needs to be passed in if there is greater than 1 unet, so that unconditional training of a single ddpm is seamless (experiment in progress locally)

dalle2_pytorch/__init__.py

@@ -1,6 +1,6 @@
 from dalle2_pytorch.dalle2_pytorch import DALLE2, DiffusionPriorNetwork, DiffusionPrior, Unet, Decoder
 from dalle2_pytorch.dalle2_pytorch import OpenAIClipAdapter
-from dalle2_pytorch.train import DecoderTrainer, DiffusionPriorTrainer
+from dalle2_pytorch.trainer import DecoderTrainer, DiffusionPriorTrainer
 
 from dalle2_pytorch.vqgan_vae import VQGanVAE
 from x_clip import CLIP

dalle2_pytorch/trainer.py

@@ -377,8 +377,11 @@ class DecoderTrainer(nn.Module):
         scaler = getattr(self, f'scaler{index}')
         return scaler.scale(loss)
 
-    def update(self, unet_number):
-        assert 1 <= unet_number <= self.num_unets
+    def update(self, unet_number = None):
+        if self.num_unets == 1:
+            unet_number = default(unet_number, 1)
+
+        assert exists(unet_number) and 1 <= unet_number <= self.num_unets
         index = unet_number - 1
         unet = self.decoder.unets[index]
 
@@ -421,10 +424,13 @@ class DecoderTrainer(nn.Module):
     def forward(
         self,
         *args,
-        unet_number,
+        unet_number = None,
         max_batch_size = None,
         **kwargs
     ):
+        if self.num_unets == 1:
+            unet_number = default(unet_number, 1)
+
         total_loss = 0.
 
         for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs):

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.2.38',
+  version = '0.2.39',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',

train_diffusion_prior.py

@@ -7,9 +7,11 @@ import numpy as np
 import torch
 import clip
 from torch import nn
+
 from dalle2_pytorch.dataloaders import make_splits
 from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork, OpenAIClipAdapter
-from dalle2_pytorch.train import DiffusionPriorTrainer, load_diffusion_model, save_diffusion_model, print_ribbon
+from dalle2_pytorch.trainer import DiffusionPriorTrainer, load_diffusion_model, save_diffusion_model, print_ribbon
+
 from dalle2_pytorch.trackers import ConsoleTracker, WandbTracker
 
 from embedding_reader import EmbeddingReader