code: import argparse
  from random import choice
  from pathlib import Path
  
  # torch
  
  import torch
  from torch.optim import Adam
  from torch.nn.utils import clip_grad_norm_
  
  # vision imports
  
  from PIL import Image
  from torchvision import transforms as T
  from torch.utils.data import DataLoader, Dataset
  from torchvision.datasets import ImageFolder
  from torchvision.utils import make_grid, save_image
  
  # dalle related classes and utils
  
  from dalle_pytorch import OpenAIDiscreteVAE, DiscreteVAE, DALLE
  from dalle_pytorch.simple_tokenizer import tokenize, tokenizer, VOCAB_SIZE
  
  # argument parsing
  
  parser = argparse.ArgumentParser()
  
  group = parser.add_mutually_exclusive_group(required = False)
  
  group.add_argument('--vae_path', type = str,
                      help='path to your trained discrete VAE')
  
  group.add_argument('--dalle_path', type = str,
                      help='path to your partially trained DALL-E')
  
  parser.add_argument('--image_text_folder', type = str, required = True,
                      help='path to your folder of images and text for learning the DALL-E')
  
  args = parser.parse_args()
  
  # helpers
  
  def exists(val):
      return val is not None
  
  # constants
  
  VAE_PATH = args.vae_path
  DALLE_PATH = args.dalle_path
  RESUME = exists(DALLE_PATH)
  
  EPOCHS = 20
  BATCH_SIZE = 4
  LEARNING_RATE = 3e-4
  GRAD_CLIP_NORM = 0.5
  
  MODEL_DIM = 512
  TEXT_SEQ_LEN = 256
  DEPTH = 2
  HEADS = 4
  DIM_HEAD = 64
  
  # reconstitute vae
  
  if RESUME:
      dalle_path = Path(DALLE_PATH)
      assert dalle_path.exists(), 'DALL-E model file does not exist'
  
      loaded_obj = torch.load(str(dalle_path))
  
      dalle_params, vae_params, weights = loaded_obj['hparams'], loaded_obj['vae_params'], loaded_obj['weights']
  
      vae = DiscreteVAE(**vae_params)
  
      dalle_params = dict(
          vae = vae,
          **dalle_params
      )
  
      IMAGE_SIZE = vae_params['image_size']
  
  else:
      if exists(VAE_PATH):
          vae_path = Path(VAE_PATH)
          assert vae_path.exists(), 'VAE model file does not exist'
  
          loaded_obj = torch.load(str(vae_path))
  
          vae_params, weights = loaded_obj['hparams'], loaded_obj['weights']
  
          vae = DiscreteVAE(**vae_params)
          vae.load_state_dict(weights)
      else:
          print('using OpenAIs pretrained VAE for encoding images to tokens')
          vae_params = None
  
          vae = OpenAIDiscreteVAE()
  
      IMAGE_SIZE = vae.image_size
  
      dalle_params = dict(
          vae = vae,
          num_text_tokens = VOCAB_SIZE,
          text_seq_len = TEXT_SEQ_LEN,
          dim = MODEL_DIM,
          depth = DEPTH,
          heads = HEADS,
          dim_head = DIM_HEAD
      )
  
  # helpers
  
  def save_model(path):
      save_obj = {
          'hparams': dalle_params,
          'vae_params': vae_params,
          'weights': dalle.state_dict()
      }
  
      torch.save(save_obj, path)
  
  # dataset loading
  
  class TextImageDataset(Dataset):
      def __init__(self, folder, text_len = 256, image_size = 128):
          super().__init__()
          path = Path(folder)
  
          text_files = [*path.glob('**/*.txt')]
  
          image_files = [
              *path.glob('**/*.png'),
              *path.glob('**/*.jpg'),
              *path.glob('**/*.jpeg')
          ]
  
          text_files = {t.stem: t for t in text_files}
          image_files = {i.stem: i for i in image_files}
  
          keys = (image_files.keys() & text_files.keys())
  
          self.keys = list(keys)
          self.text_files = {k: v for k, v in text_files.items() if k in keys}
          self.image_files = {k: v for k, v in image_files.items() if k in keys}
  
          self.image_tranform = T.Compose([
              T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
              T.CenterCrop(image_size),
              T.Resize(image_size),
              T.ToTensor(),
              T.Lambda(lambda t: t.expand(3, -1, -1)),
              T.Normalize((0.5,) * 3, (0.5,) * 3)
          ])
  
      def __len__(self):
          return len(self.keys)
  
      def __getitem__(self, ind):
          key = self.keys[ind]
          text_file = self.text_files[key]
          image_file = self.image_files[key]
  
          image = Image.open(image_file)
          descriptions = text_file.read_text().split('\n')
          descriptions = list(filter(lambda t: len(t) > 0, descriptions))
          description = choice(descriptions)
  
          tokenized_text = tokenize(description).squeeze(0)
          mask = tokenized_text != 0
  
          image_tensor = self.image_tranform(image)
          return tokenized_text, image_tensor, mask
  
  # create dataset and dataloader
  
  ds = TextImageDataset(
      args.image_text_folder,
      text_len = TEXT_SEQ_LEN,
      image_size = IMAGE_SIZE
  )
  
  assert len(ds) > 0, 'dataset is empty'
  print(f'{len(ds)} image-text pairs found for training')
  
  dl = DataLoader(ds, batch_size = BATCH_SIZE, shuffle = True, drop_last = True)
  
  # initialize DALL-E
  
  dalle = DALLE(**dalle_params).cuda()
  
  if RESUME:
      dalle.load_state_dict(weights)
  
  # optimizer
  
  opt = Adam(dalle.parameters(), lr = LEARNING_RATE)
  
  # experiment tracker
  
  import wandb
  
  wandb.config.depth = DEPTH
  wandb.config.heads = HEADS
  wandb.config.dim_head = DIM_HEAD
  
  wandb.init(project = 'dalle_train_transformer', resume = RESUME)
  
  # training
  
  for epoch in range(EPOCHS):
      for i, (text, images, mask) in enumerate(dl):
          text, images, mask = map(lambda t: t.cuda(), (text, images, mask))
  
          loss = dalle(text, images, mask = mask, return_loss = True)
  
          loss.backward()
          clip_grad_norm_(dalle.parameters(), GRAD_CLIP_NORM)
  
          opt.step()
          opt.zero_grad()
  
          log = {}
  
          if i % 10 == 0:
              print(epoch, i, f'loss - {loss.item()}')
  
              log = {
                  **log,
                  'epoch': epoch,
                  'iter': i,
                  'loss': loss.item()
              }
  
          if i % 100 == 0:
              sample_text = text[:1]
              token_list = sample_text.masked_select(sample_text != 0).tolist()
              decoded_text = tokenizer.decode(token_list)
  
              image = dalle.generate_images(
                  text[:1],
                  mask = mask[:1],
                  filter_thres = 0.9    # topk sampling at 0.9
              )
  
              save_model(f'./dalle.pt')
              wandb.save(f'./dalle.pt')
  
              log = {
                  **log,
                  'image': wandb.Image(image, caption = decoded_text)
              }
  
          wandb.log(log)
  
  save_model(f'./dalle-final.pt')
  wandb.save('./dalle-final.pt')
  wandb.finish()
  
apis: torch.utils.data.DataLoader torch.save