complete ddim integration of diffusion prior as well as decoder for each unet, feature complete for https://github.com/lucidrains/DALLE2-pytorch/issues/157

.github/workflows/ci.yml

@@ -0,0 +1,33 @@
+name: Continuous integration
+
+on:
+  push:
+    branches:
+    - main
+  pull_request:
+    branches:
+    - main
+
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.8]
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install
+      run: |
+        python3 -m venv .env
+        source .env/bin/activate
+        make install
+    - name: Tests
+      run: |
+        source .env/bin/activate
+        make test
+

.gitignore

@@ -136,3 +136,5 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+.tracker_data
+*.pth

Makefile

@@ -0,0 +1,6 @@
+install:
+	pip install -U pip
+	pip install -e .
+
+test:
+	CUDA_VISIBLE_DEVICES= python train_decoder.py --config_file configs/train_decoder_config.test.json

README.md

@@ -583,6 +583,7 @@ unet1 = Unet(
     cond_dim = 128,
     channels = 3,
     dim_mults=(1, 2, 4, 8),
+    text_embed_dim = 512,
     cond_on_text_encodings = True  # set to True for any unets that need to be conditioned on text encodings (ex. first unet in cascade)
 ).cuda()
 
@@ -598,7 +599,8 @@ decoder = Decoder(
     unet = (unet1, unet2),
     image_sizes = (128, 256),
     clip = clip,
-    timesteps = 100,
+    timesteps = 1000,
+    sample_timesteps = (250, 27),
     image_cond_drop_prob = 0.1,
     text_cond_drop_prob = 0.5
 ).cuda()

configs/train_decoder_config.test.json

@@ -0,0 +1,102 @@
+{
+    "decoder": {
+        "unets": [
+            {
+                "dim": 16,
+                "image_embed_dim": 768,
+                "cond_dim": 16,
+                "channels": 3,
+                "dim_mults": [1, 2, 4, 8],
+                "attn_dim_head": 16,
+                "attn_heads": 4,
+		"self_attn": [false, true, true, true]
+            }
+        ],
+        "clip": {
+            "make": "openai",
+            "model": "ViT-L/14"
+        },
+
+	"timesteps": 10,
+        "image_sizes": [64],
+        "channels": 3,
+        "loss_type": "l2",
+        "beta_schedule": ["cosine"],
+        "learned_variance": true
+    },
+    "data": {
+        "webdataset_base_url": "test_data/{}.tar",
+        "num_workers": 4,
+        "batch_size": 4,
+        "start_shard": 0,
+        "end_shard": 9,
+        "shard_width": 1,
+        "index_width": 1,
+        "splits": {
+            "train": 0.75,
+            "val": 0.15,
+            "test": 0.1
+        },
+        "shuffle_train": false,
+        "resample_train": true,
+        "preprocessing": {
+            "RandomResizedCrop": {
+                "size": [224, 224],
+                "scale": [0.75, 1.0],
+                "ratio": [1.0, 1.0]
+            },
+            "ToTensor": true
+        }
+    },
+    "train": {
+        "epochs": 1,
+        "lr": 1e-16,
+        "wd": 0.01,
+        "max_grad_norm": 0.5,
+        "save_every_n_samples": 100,
+        "n_sample_images": 1,
+        "device": "cpu",
+        "epoch_samples": 50,
+        "validation_samples": 5,
+        "use_ema": true,
+        "ema_beta": 0.99,
+        "amp": false,
+        "save_all": false,
+        "save_latest": true,
+        "save_best": true,
+        "unet_training_mask": [true]
+    },
+    "evaluate": {
+        "n_evaluation_samples": 2,
+        "FID": {
+            "feature": 64
+        },
+        "IS": {
+            "feature": 64,
+            "splits": 10
+        },
+        "KID": {
+            "feature": 64,
+            "subset_size": 2
+        },
+        "LPIPS": {
+            "net_type": "vgg",
+            "reduction": "mean"
+        }
+    },
+    "tracker": {
+        "overwrite_data_path": true,
+
+	"log": {
+            "log_type": "console"
+	},
+
+        "load": {
+            "load_from": null
+        },
+
+       "save": [{
+            "save_to": "local"
+        }]
+    }
+}

dalle2_pytorch/dalle2_pytorch.py

@@ -169,6 +169,11 @@ class BaseClipAdapter(nn.Module):
         self.clip = clip
         self.overrides = kwargs
 
+    def validate_and_resize_image(self, image):
+        image_size = image.shape[-1]
+        assert image_size >= self.image_size, f'you are passing in an image of size {image_size} but CLIP requires the image size to be at least {self.image_size}'
+        return resize_image_to(image, self.image_size)
+
     @property
     def dim_latent(self):
         raise NotImplementedError
@@ -219,7 +224,7 @@ class XClipAdapter(BaseClipAdapter):
 
     @torch.no_grad()
     def embed_image(self, image):
-        image = resize_image_to(image, self.image_size)
+        image = self.validate_and_resize_image(image)
         encoder_output = self.clip.visual_transformer(image)
         image_cls, image_encodings = encoder_output[:, 0], encoder_output[:, 1:]
         image_embed = self.clip.to_visual_latent(image_cls)
@@ -254,7 +259,7 @@ class CoCaAdapter(BaseClipAdapter):
 
     @torch.no_grad()
     def embed_image(self, image):
-        image = resize_image_to(image, self.image_size)
+        image = self.validate_and_resize_image(image)
         image_embed, image_encodings = self.clip.embed_image(image)
         return EmbeddedImage(image_embed, image_encodings)
 
@@ -315,7 +320,7 @@ class OpenAIClipAdapter(BaseClipAdapter):
     @torch.no_grad()
     def embed_image(self, image):
         assert not self.cleared
-        image = resize_image_to(image, self.image_size)
+        image = self.validate_and_resize_image(image)
         image = self.clip_normalize(image)
         image_embed = self.clip.encode_image(image)
         return EmbeddedImage(l2norm(image_embed.float()), None)
@@ -500,6 +505,12 @@ class NoiseScheduler(nn.Module):
             extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
         )
 
+    def predict_noise_from_start(self, x_t, t, x0):
+        return (
+            (x0 - extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t) / \
+            extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+        )
+
     def p2_reweigh_loss(self, loss, times):
         if not self.has_p2_loss_reweighting:
             return loss
@@ -906,6 +917,7 @@ class DiffusionPrior(nn.Module):
         image_size = None,
         image_channels = 3,
         timesteps = 1000,
+        sample_timesteps = None,
         cond_drop_prob = 0.,
         loss_type = "l2",
         predict_x_start = True,
@@ -919,6 +931,8 @@ class DiffusionPrior(nn.Module):
     ):
         super().__init__()
 
+        self.sample_timesteps = sample_timesteps
+
         self.noise_scheduler = NoiseScheduler(
             beta_schedule = beta_schedule,
             timesteps = timesteps,
@@ -973,8 +987,6 @@ class DiffusionPrior(nn.Module):
 
         if self.predict_x_start:
             x_recon = pred
-            # not 100% sure of this above line - for any spectators, let me know in the github issues (or through a pull request) if you know how to correctly do this
-            # i'll be rereading https://arxiv.org/abs/2111.14822, where i think a similar approach is taken
         else:
             x_recon = self.noise_scheduler.predict_start_from_noise(x, t = t, noise = pred)
 
@@ -997,21 +1009,75 @@ class DiffusionPrior(nn.Module):
         return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
 
     @torch.no_grad()
-    def p_sample_loop(self, shape, text_cond, cond_scale = 1.):
-        device = self.device
-
-        b = shape[0]
-        image_embed = torch.randn(shape, device=device)
+    def p_sample_loop_ddpm(self, shape, text_cond, cond_scale = 1.):
+        batch, device = shape[0], self.device
+        image_embed = torch.randn(shape, device = device)
 
         if self.init_image_embed_l2norm:
             image_embed = l2norm(image_embed) * self.image_embed_scale
 
         for i in tqdm(reversed(range(0, self.noise_scheduler.num_timesteps)), desc='sampling loop time step', total=self.noise_scheduler.num_timesteps):
-            times = torch.full((b,), i, device = device, dtype = torch.long)
+            times = torch.full((batch,), i, device = device, dtype = torch.long)
             image_embed = self.p_sample(image_embed, times, text_cond = text_cond, cond_scale = cond_scale)
 
         return image_embed
 
+    @torch.no_grad()
+    def p_sample_loop_ddim(self, shape, text_cond, *, timesteps, eta = 1., cond_scale = 1.):
+        batch, device, alphas, total_timesteps = shape[0], self.device, self.noise_scheduler.alphas_cumprod_prev, self.noise_scheduler.num_timesteps
+
+        times = torch.linspace(0., total_timesteps, steps = timesteps + 2)[:-1]
+
+        times = list(reversed(times.int().tolist()))
+        time_pairs = list(zip(times[:-1], times[1:]))
+
+        image_embed = torch.randn(shape, device = device)
+
+        if self.init_image_embed_l2norm:
+            image_embed = l2norm(image_embed) * self.image_embed_scale
+
+        for time, time_next in tqdm(time_pairs, desc = 'sampling loop time step'):
+            alpha = alphas[time]
+            alpha_next = alphas[time_next]
+
+            time_cond = torch.full((batch,), time, device = device, dtype = torch.long)
+
+            pred = self.net.forward_with_cond_scale(image_embed, time_cond, cond_scale = cond_scale, **text_cond)
+
+            if self.predict_x_start:
+                x_start = pred
+                pred_noise = self.noise_scheduler.predict_noise_from_start(image_embed, t = time_cond, x0 = pred)
+            else:
+                x_start = self.noise_scheduler.predict_start_from_noise(image_embed, t = time_cond, noise = pred)
+                pred_noise = pred
+
+            if not self.predict_x_start:
+                x_start.clamp_(-1., 1.)
+
+            if self.predict_x_start and self.sampling_clamp_l2norm:
+                x_start = l2norm(x_start) * self.image_embed_scale
+
+            c1 = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
+            c2 = ((1 - alpha_next) - torch.square(c1)).sqrt()
+            new_noise = torch.randn_like(image_embed)
+
+            img = x_start * alpha_next.sqrt() + \
+                  c1 * new_noise + \
+                  c2 * pred_noise
+
+        return image_embed
+
+    @torch.no_grad()
+    def p_sample_loop(self, *args, timesteps = None, **kwargs):
+        timesteps = default(timesteps, self.noise_scheduler.num_timesteps)
+        assert timesteps <= self.noise_scheduler.num_timesteps
+        is_ddim = timesteps < self.noise_scheduler.num_timesteps
+
+        if not is_ddim:
+            return self.p_sample_loop_ddpm(*args, **kwargs)
+
+        return self.p_sample_loop_ddim(*args, **kwargs, timesteps = timesteps)
+
     def p_losses(self, image_embed, times, text_cond, noise = None):
         noise = default(noise, lambda: torch.randn_like(image_embed))
 
@@ -1046,7 +1112,15 @@ class DiffusionPrior(nn.Module):
 
     @torch.no_grad()
     @eval_decorator
-    def sample(self, text, num_samples_per_batch = 2, cond_scale = 1.):
+    def sample(
+        self,
+        text,
+        num_samples_per_batch = 2,
+        cond_scale = 1.,
+        timesteps = None
+    ):
+        timesteps = default(timesteps, self.sample_timesteps)
+
         # in the paper, what they did was
         # sample 2 image embeddings, choose the top 1 similarity, as judged by CLIP
         text = repeat(text, 'b ... -> (b r) ...', r = num_samples_per_batch)
@@ -1061,7 +1135,7 @@ class DiffusionPrior(nn.Module):
         if self.condition_on_text_encodings:
             text_cond = {**text_cond, 'text_encodings': text_encodings, 'mask': text_mask}
 
-        image_embeds = self.p_sample_loop((batch_size, image_embed_dim), text_cond = text_cond, cond_scale = cond_scale)
+        image_embeds = self.p_sample_loop((batch_size, image_embed_dim), text_cond = text_cond, cond_scale = cond_scale, timesteps = timesteps)
 
         # retrieve original unscaled image embed
 
@@ -1848,6 +1922,7 @@ class Decoder(nn.Module):
         channels = 3,
         vae = tuple(),
         timesteps = 1000,
+        sample_timesteps = None,
         image_cond_drop_prob = 0.1,
         text_cond_drop_prob = 0.5,
         loss_type = 'l2',
@@ -1871,7 +1946,8 @@ class Decoder(nn.Module):
         use_dynamic_thres = False,                  # from the Imagen paper
         dynamic_thres_percentile = 0.9,
         p2_loss_weight_gamma = 0.,                  # p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time - 1. is recommended
-        p2_loss_weight_k = 1
+        p2_loss_weight_k = 1,
+        ddim_sampling_eta = 1.                      # can be set to 0. for deterministic sampling afaict
     ):
         super().__init__()
 
@@ -1951,6 +2027,11 @@ class Decoder(nn.Module):
             self.unets.append(one_unet)
             self.vaes.append(one_vae.copy_for_eval())
 
+        # sampling timesteps, defaults to non-ddim with full timesteps sampling
+
+        self.sample_timesteps = cast_tuple(sample_timesteps, num_unets)
+        self.ddim_sampling_eta = ddim_sampling_eta
+
         # create noise schedulers per unet
 
         if not exists(beta_schedule):
@@ -1961,7 +2042,9 @@ class Decoder(nn.Module):
 
         self.noise_schedulers = nn.ModuleList([])
 
-        for unet_beta_schedule, unet_p2_loss_weight_gamma in zip(beta_schedule, p2_loss_weight_gamma):
+        for ind, (unet_beta_schedule, unet_p2_loss_weight_gamma, sample_timesteps) in enumerate(zip(beta_schedule, p2_loss_weight_gamma, self.sample_timesteps)):
+            assert sample_timesteps <= timesteps, f'sampling timesteps {sample_timesteps} must be less than or equal to the number of training timesteps {timesteps} for unet {ind + 1}'
+
             noise_scheduler = NoiseScheduler(
                 beta_schedule = unet_beta_schedule,
                 timesteps = timesteps,
@@ -2062,6 +2145,26 @@ class Decoder(nn.Module):
         for unet, device in zip(self.unets, devices):
             unet.to(device)
 
+    def dynamic_threshold(self, x):
+        """ proposed in https://arxiv.org/abs/2205.11487 as an improved clamping in the setting of classifier free guidance """
+        
+        # s is the threshold amount
+        # static thresholding would just be s = 1
+        s = 1.
+        if self.use_dynamic_thres:
+            s = torch.quantile(
+                rearrange(x, 'b ... -> b (...)').abs(),
+                self.dynamic_thres_percentile,
+                dim = -1
+            )
+
+            s.clamp_(min = 1.)
+            s = s.view(-1, *((1,) * (x.ndim - 1)))
+
+        # clip by threshold, depending on whether static or dynamic
+        x = x.clamp(-s, s) / s
+        return x
+
     def p_mean_variance(self, unet, x, t, image_embed, noise_scheduler, text_encodings = None, text_mask = None, lowres_cond_img = None, clip_denoised = True, predict_x_start = False, learned_variance = False, cond_scale = 1., model_output = None):
         assert not (cond_scale != 1. and not self.can_classifier_guidance), 'the decoder was not trained with conditional dropout, and thus one cannot use classifier free guidance (cond_scale anything other than 1)'
 
@@ -2076,21 +2179,7 @@ class Decoder(nn.Module):
             x_recon = noise_scheduler.predict_start_from_noise(x, t = t, noise = pred)
 
         if clip_denoised:
-            # s is the threshold amount
-            # static thresholding would just be s = 1
-            s = 1.
-            if self.use_dynamic_thres:
-                s = torch.quantile(
-                    rearrange(x_recon, 'b ... -> b (...)').abs(),
-                    self.dynamic_thres_percentile,
-                    dim = -1
-                )
-
-                s.clamp_(min = 1.)
-                s = s.view(-1, *((1,) * (x_recon.ndim - 1)))
-
-            # clip by threshold, depending on whether static or dynamic
-            x_recon = x_recon.clamp(-s, s) / s
+            x_recon = self.dynamic_threshold(x_recon)
 
         model_mean, posterior_variance, posterior_log_variance = noise_scheduler.q_posterior(x_start=x_recon, x_t=x, t=t)
 
@@ -2120,7 +2209,7 @@ class Decoder(nn.Module):
         return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
 
     @torch.no_grad()
-    def p_sample_loop(self, unet, shape, image_embed, noise_scheduler, predict_x_start = False, learned_variance = False, clip_denoised = True, lowres_cond_img = None, text_encodings = None, text_mask = None, cond_scale = 1, is_latent_diffusion = False):
+    def p_sample_loop_ddpm(self, unet, shape, image_embed, noise_scheduler, predict_x_start = False, learned_variance = False, clip_denoised = True, lowres_cond_img = None, text_encodings = None, text_mask = None, cond_scale = 1, is_latent_diffusion = False):
         device = self.device
 
         b = shape[0]
@@ -2148,6 +2237,61 @@ class Decoder(nn.Module):
         unnormalize_img = self.unnormalize_img(img)
         return unnormalize_img
 
+    @torch.no_grad()
+    def p_sample_loop_ddim(self, unet, shape, image_embed, noise_scheduler, timesteps, eta = 1., predict_x_start = False, learned_variance = False, clip_denoised = True, lowres_cond_img = None, text_encodings = None, text_mask = None, cond_scale = 1, is_latent_diffusion = False):
+        batch, device, total_timesteps, alphas, eta = shape[0], self.device, noise_scheduler.num_timesteps, noise_scheduler.alphas_cumprod_prev, self.ddim_sampling_eta
+
+        times = torch.linspace(0., total_timesteps, steps = timesteps + 2)[:-1]
+
+        times = list(reversed(times.int().tolist()))
+        time_pairs = list(zip(times[:-1], times[1:]))
+
+        img = torch.randn(shape, device = device)
+
+        for time, time_next in tqdm(time_pairs, desc = 'sampling loop time step'):
+            alpha = alphas[time]
+            alpha_next = alphas[time_next]
+
+            time_cond = torch.full((batch,), time, device = device, dtype = torch.long)
+
+            pred = unet.forward_with_cond_scale(img, time_cond, image_embed = image_embed, text_encodings = text_encodings, text_mask = text_mask, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img)
+
+            if learned_variance:
+                pred, _ = pred.chunk(2, dim = 1)
+
+            if predict_x_start:
+                x_start = pred
+                pred_noise = noise_scheduler.predict_noise_from_start(img, t = time_cond, x0 = pred)
+            else:
+                x_start = noise_scheduler.predict_start_from_noise(img, t = time_cond, noise = pred)
+                pred_noise = pred
+
+            if clip_denoised:
+                x_start = self.dynamic_threshold(x_start)
+
+            c1 = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
+            c2 = ((1 - alpha_next) - torch.square(c1)).sqrt()
+
+            img = x_start * alpha_next.sqrt() + \
+                  c1 * torch.randn_like(img) + \
+                  c2 * pred_noise
+
+        img = self.unnormalize_img(img)
+        return img
+
+    @torch.no_grad()
+    def p_sample_loop(self, *args, noise_scheduler, timesteps = None, **kwargs):
+        num_timesteps = noise_scheduler.num_timesteps
+
+        timesteps = default(timesteps, num_timesteps)
+        assert timesteps <= num_timesteps
+        is_ddim = timesteps < num_timesteps
+
+        if not is_ddim:
+            return self.p_sample_loop_ddpm(*args, noise_scheduler = noise_scheduler, **kwargs)
+
+        return self.p_sample_loop_ddim(*args, noise_scheduler = noise_scheduler, timesteps = timesteps, **kwargs)
+
     def p_losses(self, unet, x_start, times, *, image_embed, noise_scheduler, lowres_cond_img = None, text_encodings = None, text_mask = None, predict_x_start = False, noise = None, learned_variance = False, clip_denoised = False, is_latent_diffusion = False):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
@@ -2248,7 +2392,7 @@ class Decoder(nn.Module):
         img = None
         is_cuda = next(self.parameters()).is_cuda
 
-        for unet_number, unet, vae, channel, image_size, predict_x_start, learned_variance, noise_scheduler in tqdm(zip(range(1, len(self.unets) + 1), self.unets, self.vaes, self.sample_channels, self.image_sizes, self.predict_x_start, self.learned_variance, self.noise_schedulers)):
+        for unet_number, unet, vae, channel, image_size, predict_x_start, learned_variance, noise_scheduler, sample_timesteps in tqdm(zip(range(1, len(self.unets) + 1), self.unets, self.vaes, self.sample_channels, self.image_sizes, self.predict_x_start, self.learned_variance, self.noise_schedulers, self.sample_timesteps)):
 
             context = self.one_unet_in_gpu(unet = unet) if is_cuda and not distributed else null_context()
 
@@ -2277,7 +2421,8 @@ class Decoder(nn.Module):
                     clip_denoised = not is_latent_diffusion,
                     lowres_cond_img = lowres_cond_img,
                     is_latent_diffusion = is_latent_diffusion,
-                    noise_scheduler = noise_scheduler
+                    noise_scheduler = noise_scheduler,
+                    timesteps = sample_timesteps
                 )
 
                 img = vae.decode(img)

dalle2_pytorch/dataloaders/decoder_loader.py

@@ -1,6 +1,7 @@
 import os
 import webdataset as wds
 import torch
+from torch.utils.data import DataLoader
 import numpy as np
 import fsspec
 import shutil
@@ -255,7 +256,7 @@ def create_image_embedding_dataloader(
     )
     if shuffle_num is not None and shuffle_num > 0:
         ds.shuffle(1000)
-    return wds.WebLoader(
+    return DataLoader(
         ds,
         num_workers=num_workers,
         batch_size=batch_size,

dalle2_pytorch/train_configs.py

@@ -154,6 +154,7 @@ class DiffusionPriorConfig(BaseModel):
     image_size: int
     image_channels: int = 3
     timesteps: int = 1000
+    sample_timesteps: Optional[int] = None
     cond_drop_prob: float = 0.
     loss_type: str = 'l2'
     predict_x_start: bool = True
@@ -233,6 +234,7 @@ class DecoderConfig(BaseModel):
     clip: Optional[AdapterConfig]   # The clip model to use if embeddings are not provided
     channels: int = 3
     timesteps: int = 1000
+    sample_timesteps: Optional[SingularOrIterable(int)] = None
     loss_type: str = 'l2'
     beta_schedule: ListOrTuple(str) = 'cosine'
     learned_variance: bool = True

dalle2_pytorch/trainer.py

@@ -21,7 +21,7 @@ import pytorch_warmup as warmup
 
 from ema_pytorch import EMA
 
-from accelerate import Accelerator
+from accelerate import Accelerator, DistributedType
 
 import numpy as np
 
@@ -76,6 +76,7 @@ def cast_torch_tensor(fn):
     def inner(model, *args, **kwargs):
         device = kwargs.pop('_device', next(model.parameters()).device)
         cast_device = kwargs.pop('_cast_device', True)
+        cast_deepspeed_precision = kwargs.pop('_cast_deepspeed_precision', True)
 
         kwargs_keys = kwargs.keys()
         all_args = (*args, *kwargs.values())
@@ -85,6 +86,21 @@ def cast_torch_tensor(fn):
         if cast_device:
             all_args = tuple(map(lambda t: t.to(device) if exists(t) and isinstance(t, torch.Tensor) else t, all_args))
 
+        if cast_deepspeed_precision:
+            try:
+                accelerator = model.accelerator
+                if accelerator is not None and accelerator.distributed_type == DistributedType.DEEPSPEED:
+                    cast_type_map = {
+                        "fp16": torch.half,
+                        "bf16": torch.bfloat16,
+                        "no": torch.float
+                    }
+                    precision_type = cast_type_map[accelerator.mixed_precision]
+                    all_args = tuple(map(lambda t: t.to(precision_type) if exists(t) and isinstance(t, torch.Tensor) else t, all_args))
+            except AttributeError:
+                # Then this model doesn't have an accelerator
+                pass
+
         args, kwargs_values = all_args[:split_kwargs_index], all_args[split_kwargs_index:]
         kwargs = dict(tuple(zip(kwargs_keys, kwargs_values)))
 
@@ -446,6 +462,7 @@ class DecoderTrainer(nn.Module):
         self,
         decoder,
         accelerator = None,
+        dataloaders = None,
         use_ema = True,
         lr = 1e-4,
         wd = 1e-2,
@@ -508,10 +525,31 @@ class DecoderTrainer(nn.Module):
 
         self.register_buffer('steps', torch.tensor([0] * self.num_unets))
 
+        if self.accelerator.distributed_type == DistributedType.DEEPSPEED and decoder.clip is not None:
+            # Then we need to make sure clip is using the correct precision or else deepspeed will error
+            cast_type_map = {
+                "fp16": torch.half,
+                "bf16": torch.bfloat16,
+                "no": torch.float
+            }
+            precision_type = cast_type_map[accelerator.mixed_precision]
+            assert precision_type == torch.float, "DeepSpeed currently only supports float32 precision when using on the fly embedding generation from clip"
+            clip = decoder.clip
+            clip.to(precision_type)
+
         decoder, *optimizers = list(self.accelerator.prepare(decoder, *optimizers))
 
         self.decoder = decoder
 
+        # prepare dataloaders
+
+        train_loader = val_loader = None
+        if exists(dataloaders):
+            train_loader, val_loader = self.accelerator.prepare(dataloaders["train"], dataloaders["val"])
+
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+
         # store optimizers
 
         for opt_ind, optimizer in zip(range(len(optimizers)), optimizers):
@@ -675,6 +713,9 @@ class DecoderTrainer(nn.Module):
 
         total_loss = 0.
 
+        
+        using_amp = self.accelerator.mixed_precision != 'no'
+
         for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs):
             with self.accelerator.autocast():
                 loss = self.decoder(*chunked_args, unet_number = unet_number, **chunked_kwargs)

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.17.0'
+__version__ = '0.19.1'

train_decoder.py

@@ -132,7 +132,7 @@ def get_example_data(dataloader, device, n=5):
             break
     return list(zip(images[:n], img_embeddings[:n], text_embeddings[:n], captions[:n]))
 
-def generate_samples(trainer, example_data, condition_on_text_encodings=False, text_prepend=""):
+def generate_samples(trainer, example_data, condition_on_text_encodings=False, text_prepend="", match_image_size=True):
     """
     Takes example data and generates images from the embeddings
     Returns three lists: real images, generated images, and captions
@@ -160,6 +160,9 @@ def generate_samples(trainer, example_data, condition_on_text_encodings=False, t
     samples = trainer.sample(**sample_params)
     generated_images = list(samples)
     captions = [text_prepend + txt for txt in txts]
+    if match_image_size:
+        generated_image_size = generated_images[0].shape[-1]
+        real_images = [resize_image_to(image, generated_image_size, clamp_range=(0, 1)) for image in real_images]
     return real_images, generated_images, captions
 
 def generate_grid_samples(trainer, examples, condition_on_text_encodings=False, text_prepend=""):
@@ -167,14 +170,6 @@ def generate_grid_samples(trainer, examples, condition_on_text_encodings=False,
     Generates samples and uses torchvision to put them in a side by side grid for easy viewing
     """
     real_images, generated_images, captions = generate_samples(trainer, examples, condition_on_text_encodings, text_prepend)
-
-    real_image_size = real_images[0].shape[-1]
-    generated_image_size = generated_images[0].shape[-1]
-
-    # training images may be larger than the generated one
-    if real_image_size > generated_image_size:
-        real_images = [resize_image_to(image, generated_image_size) for image in real_images]
-
     grid_images = [torchvision.utils.make_grid([original_image, generated_image]) for original_image, generated_image in zip(real_images, generated_images)]
     return grid_images, captions
                     
@@ -279,6 +274,7 @@ def train(
     trainer = DecoderTrainer(
         decoder=decoder,
         accelerator=accelerator,
+        dataloaders=dataloaders,
         **kwargs
     )
 
@@ -289,7 +285,6 @@ def train(
     sample = 0
     samples_seen = 0
     val_sample = 0
-    step = lambda: int(trainer.num_steps_taken(unet_number=1))
 
     if tracker.can_recall:
         start_epoch, validation_losses, next_task, recalled_sample, samples_seen = recall_trainer(tracker, trainer)
@@ -304,6 +299,8 @@ def train(
     if not exists(unet_training_mask):
         # Then the unet mask should be true for all unets in the decoder
         unet_training_mask = [True] * trainer.num_unets
+    first_training_unet = min(index for index, mask in enumerate(unet_training_mask) if mask)
+    step = lambda: int(trainer.num_steps_taken(unet_number=first_training_unet+1))
     assert len(unet_training_mask) == trainer.num_unets, f"The unet training mask should be the same length as the number of unets in the decoder. Got {len(unet_training_mask)} and {trainer.num_unets}"
 
     accelerator.print(print_ribbon("Generating Example Data", repeat=40))
@@ -326,7 +323,7 @@ def train(
         last_snapshot = sample
 
         if next_task == 'train':
-            for i, (img, emb, txt) in enumerate(dataloaders["train"]):
+            for i, (img, emb, txt) in enumerate(trainer.train_loader):
                 # We want to count the total number of samples across all processes
                 sample_length_tensor[0] = len(img)
                 all_samples = accelerator.gather(sample_length_tensor)  # TODO: accelerator.reduce is broken when this was written. If it is fixed replace this.
@@ -419,7 +416,7 @@ def train(
             timer = Timer()
             accelerator.wait_for_everyone()
             i = 0
-            for i, (img, emb, txt) in enumerate(dataloaders["val"]):
+            for i, (img, emb, txt) in enumerate(trainer.val_loader):  # Use the accelerate prepared loader
                 val_sample_length_tensor[0] = len(img)
                 all_samples = accelerator.gather(val_sample_length_tensor)
                 total_samples = all_samples.sum().item()
@@ -524,6 +521,20 @@ def initialize_training(config: TrainDecoderConfig, config_path):
     # Set up accelerator for configurable distributed training
     ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=config.train.find_unused_parameters)
     accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
+
+    if accelerator.num_processes > 1:
+        # We are using distributed training and want to immediately ensure all can connect
+        accelerator.print("Waiting for all processes to connect...")
+        accelerator.wait_for_everyone()
+        accelerator.print("All processes online and connected")
+
+    # If we are in deepspeed fp16 mode, we must ensure learned variance is off
+    if accelerator.mixed_precision == "fp16" and accelerator.distributed_type == accelerate_dataclasses.DistributedType.DEEPSPEED and config.decoder.learned_variance:
+        raise ValueError("DeepSpeed fp16 mode does not support learned variance")
+
+    if accelerator.process_index != accelerator.local_process_index and accelerator.distributed_type == accelerate_dataclasses.DistributedType.DEEPSPEED:
+        # This is an invalid configuration until we figure out how to handle this
+        raise ValueError("DeepSpeed does not support multi-node distributed training")
     
     # Set up data
     all_shards = list(range(config.data.start_shard, config.data.end_shard + 1))