complete inpainting ability using inpaint_image and inpaint_mask passed into sample function for decoder

Accelerate removes the ability to collate strings. Likely since it
cannot gather strings.

.github/FUNDING.yml

@@ -1 +1 @@
-github: [lucidrains]
+github: [nousr, Veldrovive, lucidrains]

.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

@@ -45,6 +45,7 @@ This library would not have gotten to this working state without the help of
 - <a href="https://github.com/rom1504">Romain</a> for the pull request reviews and project management
 - <a href="https://github.com/Ciaohe">He Cao</a> and <a href="https://github.com/xiankgx">xiankgx</a> for the Q&A and for identifying of critical bugs
 - <a href="https://github.com/marunine">Marunine</a> for identifying issues with resizing of the low resolution conditioner, when training the upsampler, in addition to various other bug fixes
+- <a href="https://github.com/malumadev">MalumaDev</a> for proposing the use of pixel shuffle upsampler for fixing checkboard artifacts
 - <a href="https://github.com/crowsonkb">Katherine</a> for her advice
 - <a href="https://stability.ai/">Stability AI</a> for the generous sponsorship
 - <a href="https://huggingface.co">🤗 Huggingface</a> and in particular <a href="https://github.com/sgugger">Sylvain</a> for the <a href="https://github.com/huggingface/accelerate">Accelerate</a> library
@@ -355,7 +356,8 @@ prior_network = DiffusionPriorNetwork(
 diffusion_prior = DiffusionPrior(
     net = prior_network,
     clip = clip,
-    timesteps = 100,
+    timesteps = 1000,
+    sample_timesteps = 64,
     cond_drop_prob = 0.2
 ).cuda()
 
@@ -419,7 +421,7 @@ For the layperson, no worries, training will all be automated into a CLI tool, a
 
 ## Training on Preprocessed CLIP Embeddings
 
-It is likely, when scaling up, that you would first preprocess your images and text into corresponding embeddings before training the prior network. You can do so easily by simply passing in `image_embed`, `text_embed`, and optionally `text_encodings` and `text_mask`
+It is likely, when scaling up, that you would first preprocess your images and text into corresponding embeddings before training the prior network. You can do so easily by simply passing in `image_embed`, `text_embed`, and optionally `text_encodings`
 
 Working example below
 
@@ -583,6 +585,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 +601,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()
@@ -1044,11 +1048,10 @@ Once built, images will be saved to the same directory the command is invoked
 - [x] bring in skip-layer excitations (from lightweight gan paper) to see if it helps for either decoder of unet or vqgan-vae training (doesnt work well)
 - [x] test out grid attention in cascading ddpm locally, decide whether to keep or remove https://arxiv.org/abs/2204.01697 (keeping, seems to be fine)
 - [x] allow for unet to be able to condition non-cross attention style as well
-- [ ] become an expert with unets, cleanup unet code, make it fully configurable, port all learnings over to https://github.com/lucidrains/x-unet (test out unet² in ddpm repo) - consider https://github.com/lucidrains/uformer-pytorch attention-based unet
-- [ ] speed up inference, read up on papers (ddim or diffusion-gan, etc)
-- [ ] figure out if possible to augment with external memory, as described in https://arxiv.org/abs/2204.11824
+- [x] speed up inference, read up on papers (ddim)
+- [x] add inpainting ability using resampler from repaint paper https://arxiv.org/abs/2201.09865
+- [ ] try out the nested unet from https://arxiv.org/abs/2005.09007 after hearing several positive testimonies from researchers, for segmentation anyhow
 - [ ] interface out the vqgan-vae so a pretrained one can be pulled off the shelf to validate latent diffusion + DALL-E2
-- [ ] add inpainting ability using resampler from repaint paper https://arxiv.org/abs/2201.09865
 
 ## Citations
 

configs/README.md

@@ -30,6 +30,7 @@ Defines the configuration options for the decoder model. The unets defined above
 | `loss_type` | No | `l2` | The loss function. Options are `l1`, `huber`, or `l2`. |
 | `beta_schedule` | No | `cosine` | The noising schedule. Options are `cosine`, `linear`, `quadratic`, `jsd`, or `sigmoid`. |
 | `learned_variance` | No | `True` | Whether to learn the variance. |
+| `clip` | No | `None` | The clip model to use if embeddings are being generated on the fly. Takes keys `make` and `model` with defaults `openai` and `ViT-L/14`. |
 
 Any parameter from the `Decoder` constructor can also be given here.
 
@@ -39,7 +40,8 @@ Settings for creation of the dataloaders.
 | Option | Required | Default | Description |
 | ------ | -------- | ------- | ----------- |
 | `webdataset_base_url` | Yes | N/A | The url of a shard in the webdataset with the shard replaced with `{}`[^1]. |
-| `embeddings_url` | No | N/A | The url of the folder containing embeddings shards. Not required if embeddings are in webdataset. |
+| `img_embeddings_url` | No | `None` | The url of the folder containing image embeddings shards. Not required if embeddings are in webdataset or clip is being used. |
+| `text_embeddings_url` | No | `None` | The url of the folder containing text embeddings shards. Not required if embeddings are in webdataset or clip is being used. |
 | `num_workers` | No | `4` | The number of workers used in the dataloader. |
 | `batch_size` | No | `64` | The batch size. |
 | `start_shard` | No | `0` | Defines the start of the shard range the dataset will recall. |
@@ -106,6 +108,13 @@ Tracking is split up into three sections:
 
 **Logging:**
 
+All loggers have the following keys:
+| Option | Required | Default | Description |
+| ------ | -------- | ------- | ----------- |
+| `log_type` | Yes | N/A | The type of logger class to use. |
+| `resume` | No | `False` | For loggers that have the option to resume an old run, resume it using maually input parameters. |
+| `auto_resume` | No | `False` | If true, the logger will attempt to resume an old run using parameters from that previous run. |
+
 If using `console` there is no further configuration than setting `log_type` to `console`.
 | Option | Required | Default | Description |
 | ------ | -------- | ------- | ----------- |
@@ -119,10 +128,15 @@ If using `wandb`
 | `wandb_project` | Yes | N/A | The wandb project save the run to. |
 | `wandb_run_name` | No | `None` | The wandb run name. |
 | `wandb_run_id` | No | `None` | The wandb run id. Used if resuming an old run. |
-| `wandb_resume` | No | `False` | Whether to resume an old run. |
 
 **Loading:**
 
+All loaders have the following keys:
+| Option | Required | Default | Description |
+| ------ | -------- | ------- | ----------- |
+| `load_from` | Yes | N/A | The type of loader class to use. |
+| `only_auto_resume` | No | `False` | If true, the loader will only load the model if the run is being auto resumed. |
+
 If using `local`
 | Option | Required | Default | Description |
 | ------ | -------- | ------- | ----------- |

configs/train_decoder_config.example.json

@@ -20,7 +20,7 @@
     },
     "data": {
         "webdataset_base_url": "pipe:s3cmd get s3://bucket/path/{}.tar -",
-        "embeddings_url": "s3://bucket/embeddings/path/",
+        "img_embeddings_url": "s3://bucket/img_embeddings/path/",
         "num_workers": 4,
         "batch_size": 64,
         "start_shard": 0,

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/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/trackers.py

@@ -1,5 +1,6 @@
 import urllib.request
 import os
+import json
 from pathlib import Path
 import shutil
 from itertools import zip_longest
@@ -37,14 +38,17 @@ class BaseLogger:
         data_path (str): A file path for storing temporary data.
         verbose (bool): Whether of not to always print logs to the console.
     """
-    def __init__(self, data_path: str, verbose: bool = False, **kwargs):
+    def __init__(self, data_path: str, resume: bool = False, auto_resume: bool = False, verbose: bool = False, **kwargs):
         self.data_path = Path(data_path)
+        self.resume = resume
+        self.auto_resume = auto_resume
         self.verbose = verbose
 
     def init(self, full_config: BaseModel, extra_config: dict, **kwargs) -> None:
         """
         Initializes the logger.
         Errors if the logger is invalid.
+        full_config is the config file dict while extra_config is anything else from the script that is not defined the config file.
         """
         raise NotImplementedError
 
@@ -60,6 +64,14 @@ class BaseLogger:
     def log_error(self, error_string, **kwargs) -> None:
         raise NotImplementedError
 
+    def get_resume_data(self, **kwargs) -> dict:
+        """
+        Sets tracker attributes that along with { "resume": True } will be used to resume training.
+        It is assumed that after init is called this data will be complete.
+        If the logger does not have any resume functionality, it should return an empty dict.
+        """
+        raise NotImplementedError
+
 class ConsoleLogger(BaseLogger):
     def init(self, full_config: BaseModel, extra_config: dict, **kwargs) -> None:
         print("Logging to console")
@@ -76,6 +88,9 @@ class ConsoleLogger(BaseLogger):
     def log_error(self, error_string, **kwargs) -> None:
         print(error_string)
 
+    def get_resume_data(self, **kwargs) -> dict:
+        return {}
+
 class WandbLogger(BaseLogger):
     """
     Logs to a wandb run.
@@ -85,7 +100,6 @@ class WandbLogger(BaseLogger):
         wandb_project (str): The wandb project to log to.
         wandb_run_id (str): The wandb run id to resume.
         wandb_run_name (str): The wandb run name to use.
-        wandb_resume (bool): Whether to resume a wandb run.
     """
     def __init__(self,
         data_path: str,
@@ -93,7 +107,6 @@ class WandbLogger(BaseLogger):
         wandb_project: str,
         wandb_run_id: Optional[str] = None,
         wandb_run_name: Optional[str] = None,
-        wandb_resume: bool = False,
         **kwargs
     ):
         super().__init__(data_path, **kwargs)
@@ -101,7 +114,6 @@ class WandbLogger(BaseLogger):
         self.project = wandb_project
         self.run_id = wandb_run_id
         self.run_name = wandb_run_name
-        self.resume = wandb_resume
 
     def init(self, full_config: BaseModel, extra_config: dict, **kwargs) -> None:
         assert self.entity is not None, "wandb_entity must be specified for wandb logger"
@@ -149,6 +161,14 @@ class WandbLogger(BaseLogger):
             print(error_string)
         self.wandb.log({"error": error_string, **kwargs}, step=step)
 
+    def get_resume_data(self, **kwargs) -> dict:
+        # In order to resume, we need wandb_entity, wandb_project, and wandb_run_id
+        return {
+            "entity": self.entity,
+            "project": self.project,
+            "run_id": self.wandb.run.id
+        }
+
 logger_type_map = {
     'console': ConsoleLogger,
     'wandb': WandbLogger,
@@ -168,8 +188,9 @@ class BaseLoader:
     Parameters:
         data_path (str): A file path for storing temporary data.
     """
-    def __init__(self, data_path: str, **kwargs):
+    def __init__(self, data_path: str, only_auto_resume: bool = False, **kwargs):
         self.data_path = Path(data_path)
+        self.only_auto_resume = only_auto_resume
 
     def init(self, logger: BaseLogger, **kwargs) -> None:
         raise NotImplementedError
@@ -304,6 +325,10 @@ class LocalSaver(BaseSaver):
     def save_file(self, local_path: str, save_path: str, **kwargs) -> None:
         # Copy the file to save_path
         save_path_file_name = Path(save_path).name
+        # Make sure parent directory exists
+        save_path_parent = Path(save_path).parent
+        if not save_path_parent.exists():
+            save_path_parent.mkdir(parents=True)
         print(f"Saving {save_path_file_name} {self.save_type} to local path {save_path}")
         shutil.copy(local_path, save_path)
 
@@ -385,11 +410,7 @@ class Tracker:
     def __init__(self, data_path: Optional[str] = DEFAULT_DATA_PATH, overwrite_data_path: bool = False, dummy_mode: bool = False):
         self.data_path = Path(data_path)
         if not dummy_mode:
-            if overwrite_data_path:
-                if self.data_path.exists():
-                    shutil.rmtree(self.data_path)
-                self.data_path.mkdir(parents=True)
-            else:
+            if not overwrite_data_path:
                 assert not self.data_path.exists(), f'Data path {self.data_path} already exists. Set overwrite_data_path to True to overwrite.'
                 if not self.data_path.exists():
                     self.data_path.mkdir(parents=True)
@@ -398,7 +419,46 @@ class Tracker:
         self.savers: List[BaseSaver]= []
         self.dummy_mode = dummy_mode
 
+    def _load_auto_resume(self) -> bool:
+        # If the file does not exist, we return False. If autoresume is enabled we print a warning so that the user can know that this is the first run.
+        if not self.auto_resume_path.exists():
+            if self.logger.auto_resume:
+                print("Auto_resume is enabled but no auto_resume.json file exists. Assuming this is the first run.")
+            return False
+
+        # Now we know that the autoresume file exists, but if we are not auto resuming we should remove it so that we don't accidentally load it next time
+        if not self.logger.auto_resume:
+            print(f'Removing auto_resume.json because auto_resume is not enabled in the config')
+            self.auto_resume_path.unlink()
+            return False
+
+        # Otherwise we read the json into a dictionary will will override parts of logger.__dict__
+        with open(self.auto_resume_path, 'r') as f:
+            auto_resume_dict = json.load(f)
+        # Check if the logger is of the same type as the autoresume save
+        if auto_resume_dict["logger_type"] != self.logger.__class__.__name__:
+            raise Exception(f'The logger type in the auto_resume file is {auto_resume_dict["logger_type"]} but the current logger is {self.logger.__class__.__name__}. Either use the original logger type, set `auto_resume` to `False`, or delete your existing tracker-data folder.')
+        # Then we are ready to override the logger with the autoresume save
+        self.logger.__dict__["resume"] = True
+        print(f"Updating {self.logger.__dict__} with {auto_resume_dict}")
+        self.logger.__dict__.update(auto_resume_dict)
+        return True
+
+    def _save_auto_resume(self):
+        # Gets the autoresume dict from the logger and adds "logger_type" to it then saves it to the auto_resume file
+        auto_resume_dict = self.logger.get_resume_data()
+        auto_resume_dict['logger_type'] = self.logger.__class__.__name__
+        with open(self.auto_resume_path, 'w') as f:
+            json.dump(auto_resume_dict, f)
+
     def init(self, full_config: BaseModel, extra_config: dict):
+        self.auto_resume_path = self.data_path / 'auto_resume.json'
+        # Check for resuming the run
+        self.did_auto_resume = self._load_auto_resume()
+        if self.did_auto_resume:
+            print(f'\n\nWARNING: RUN HAS BEEN AUTO-RESUMED WITH THE LOGGER TYPE {self.logger.__class__.__name__}.\nIf this was not your intention, stop this run and set `auto_resume` to `False` in the config.\n\n')
+            print(f"New logger config: {self.logger.__dict__}")
+        
         assert self.logger is not None, '`logger` must be set before `init` is called'
         if self.dummy_mode:
             # The only thing we need is a loader
@@ -406,12 +466,17 @@ class Tracker:
                 self.loader.init(self.logger)
             return
         assert len(self.savers) > 0, '`savers` must be set before `init` is called'
+
         self.logger.init(full_config, extra_config)
         if self.loader is not None:
             self.loader.init(self.logger)
         for saver in self.savers:
             saver.init(self.logger)
 
+        if self.logger.auto_resume:
+            # Then we need to save the autoresume file. It is assumed after logger.init is called that the logger is ready to be saved.
+            self._save_auto_resume()
+
     def add_logger(self, logger: BaseLogger):
         self.logger = logger
 
@@ -503,11 +568,16 @@ class Tracker:
                     self.logger.log_error(f'Error saving checkpoint: {e}', **kwargs)
                     print(f'Error saving checkpoint: {e}')
     
+    @property
+    def can_recall(self):
+        # Defines whether a recall can be performed.
+        return self.loader is not None and (not self.loader.only_auto_resume or self.did_auto_resume)
+    
     def recall(self):
-        if self.loader is not None:
+        if self.can_recall:
             return self.loader.recall()
         else:
-            raise ValueError('No loader specified')
+            raise ValueError('Tried to recall, but no loader was set or auto-resume was not performed.')
 
 
     

dalle2_pytorch/train_configs.py

@@ -47,6 +47,8 @@ class TrainSplitConfig(BaseModel):
 
 class TrackerLogConfig(BaseModel):
     log_type: str = 'console'
+    resume: bool = False  # For logs that are saved to unique locations, resume a previous run
+    auto_resume: bool = False  # If the process crashes and restarts, resume from the run that crashed
     verbose: bool = False
 
     class Config:
@@ -59,6 +61,7 @@ class TrackerLogConfig(BaseModel):
 
 class TrackerLoadConfig(BaseModel):
     load_from: Optional[str] = None
+    only_auto_resume: bool = False  # Only attempt to load if the logger is auto-resuming
 
     class Config:
         extra = "allow"
@@ -126,6 +129,7 @@ class AdapterConfig(BaseModel):
 class DiffusionPriorNetworkConfig(BaseModel):
     dim: int
     depth: int
+    max_text_len: int = None
     num_timesteps: int = None
     num_time_embeds: int = 1
     num_image_embeds: int = 1
@@ -133,6 +137,7 @@ class DiffusionPriorNetworkConfig(BaseModel):
     dim_head: int = 64
     heads: int = 8
     ff_mult: int = 4
+    norm_in: bool = False
     norm_out: bool = True
     attn_dropout: float = 0.
     ff_dropout: float = 0.
@@ -151,6 +156,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
@@ -219,6 +225,7 @@ class UnetConfig(BaseModel):
     self_attn: ListOrTuple(int)
     attn_dim_head: int = 32
     attn_heads: int = 16
+    init_cross_embed: bool = True
 
     class Config:
         extra = "allow"
@@ -230,6 +237,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)))
 
@@ -192,6 +208,7 @@ class DiffusionPriorTrainer(nn.Module):
             self.device = diffusion_prior_device
         else:
             self.device = accelerator.device if exists(accelerator) else device
+            diffusion_prior.to(self.device)
 
         # save model
 
@@ -445,6 +462,7 @@ class DecoderTrainer(nn.Module):
         self,
         decoder,
         accelerator = None,
+        dataloaders = None,
         use_ema = True,
         lr = 1e-4,
         wd = 1e-2,
@@ -507,11 +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))
-        schedulers = list(self.accelerator.prepare(*schedulers))
 
         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):
@@ -526,6 +564,17 @@ class DecoderTrainer(nn.Module):
 
         self.warmup_schedulers = warmup_schedulers
 
+    def validate_and_return_unet_number(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
+        return unet_number
+
+    def num_steps_taken(self, unet_number = None):
+        unet_number = self.validate_and_return_unet_number(unet_number)
+        return self.steps[unet_number - 1].item()
+
     def save(self, path, overwrite = True, **kwargs):
         path = Path(path)
         assert not (path.exists() and not overwrite)
@@ -594,10 +643,7 @@ class DecoderTrainer(nn.Module):
         self.steps += F.one_hot(unet_index_tensor, num_classes = len(self.steps))
 
     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
+        unet_number = self.validate_and_return_unet_number(unet_number)
         index = unet_number - 1
 
         optimizer = getattr(self, f'optim{index}')
@@ -627,8 +673,14 @@ class DecoderTrainer(nn.Module):
     def sample(self, *args, **kwargs):
         distributed = self.accelerator.num_processes > 1
         base_decoder = self.accelerator.unwrap_model(self.decoder)
+
+        was_training = base_decoder.training
+        base_decoder.eval()
+
         if kwargs.pop('use_non_ema', False) or not self.use_ema:
-            return base_decoder.sample(*args, **kwargs, distributed = distributed)
+            out = base_decoder.sample(*args, **kwargs, distributed = distributed)
+            base_decoder.train(was_training)
+            return out
 
         trainable_unets = self.accelerator.unwrap_model(self.decoder).unets
         base_decoder.unets = self.unets                  # swap in exponential moving averaged unets for sampling
@@ -641,6 +693,7 @@ class DecoderTrainer(nn.Module):
         for ema in self.ema_unets:
             ema.restore_ema_model_device()
 
+        base_decoder.train(was_training)
         return output
 
     @torch.no_grad()
@@ -663,11 +716,13 @@ class DecoderTrainer(nn.Module):
         max_batch_size = None,
         **kwargs
     ):
-        if self.num_unets == 1:
-            unet_number = default(unet_number, 1)
+        unet_number = self.validate_and_return_unet_number(unet_number)
 
         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.16.16'
+__version__ = '0.26.0'

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,9 +285,8 @@ def train(
     sample = 0
     samples_seen = 0
     val_sample = 0
-    step = lambda: int(trainer.step.item())
 
-    if tracker.loader is not None:
+    if tracker.can_recall:
         start_epoch, validation_losses, next_task, recalled_sample, samples_seen = recall_trainer(tracker, trainer)
         if next_task == 'train':
             sample = recalled_sample
@@ -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))
@@ -361,6 +358,7 @@ def train(
                         else:
                             # Then we need to pass the text instead
                             tokenized_texts = tokenize(txt, truncate=True)
+                            assert tokenized_texts.shape[0] == len(img), f"The number of texts ({tokenized_texts.shape[0]}) should be the same as the number of images ({len(img)})"
                             forward_params['text'] = tokenized_texts
                     loss = trainer.forward(img, **forward_params, unet_number=unet)
                     trainer.update(unet_number=unet)
@@ -419,7 +417,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(dataloaders['val']):  # 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 +522,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))
@@ -546,7 +558,7 @@ def initialize_training(config: TrainDecoderConfig, config_path):
 
     # Create the decoder model and print basic info
     decoder = config.decoder.create()
-    num_parameters = sum(p.numel() for p in decoder.parameters())
+    get_num_parameters = lambda model, only_training=False: sum(p.numel() for p in model.parameters() if (p.requires_grad or not only_training))
 
     # Create and initialize the tracker if we are the master
     tracker = create_tracker(accelerator, config, config_path, dummy = rank!=0)
@@ -575,7 +587,10 @@ def initialize_training(config: TrainDecoderConfig, config_path):
     accelerator.print(print_ribbon("Loaded Config", repeat=40))
     accelerator.print(f"Running training with {accelerator.num_processes} processes and {accelerator.distributed_type} distributed training")
     accelerator.print(f"Training using {data_source_string}. {'conditioned on text' if conditioning_on_text else 'not conditioned on text'}")
-    accelerator.print(f"Number of parameters: {num_parameters}")
+    accelerator.print(f"Number of parameters: {get_num_parameters(decoder)} total; {get_num_parameters(decoder, only_training=True)} training")
+    for i, unet in enumerate(decoder.unets):
+        accelerator.print(f"Unet {i} has {get_num_parameters(unet)} total; {get_num_parameters(unet, only_training=True)} training")
+
     train(dataloaders, decoder, accelerator,
         tracker=tracker,
         inference_device=accelerator.device,

train_diffusion_prior.py

@@ -126,9 +126,9 @@ def report_cosine_sims(
 
         # we are text conditioned, we produce an embedding from the tokenized text
         if text_conditioned:
-            text_embedding, text_encodings, text_mask = trainer.embed_text(text_data)
+            text_embedding, text_encodings = trainer.embed_text(text_data)
             text_cond = dict(
-                text_embed=text_embedding, text_encodings=text_encodings, mask=text_mask
+                text_embed=text_embedding, text_encodings=text_encodings
             )
         else:
             text_embedding = text_data
@@ -146,15 +146,12 @@ def report_cosine_sims(
 
         if text_conditioned:
             text_encodings_shuffled = text_encodings[rolled_idx]
-            text_mask_shuffled = text_mask[rolled_idx]
         else:
             text_encodings_shuffled = None
-            text_mask_shuffled = None
 
         text_cond_shuffled = dict(
             text_embed=text_embed_shuffled,
-            text_encodings=text_encodings_shuffled,
-            mask=text_mask_shuffled,
+            text_encodings=text_encodings_shuffled
         )
 
         # prepare the text embedding