added pre-commit (#5)

* added pre-commit

* polish

.isort.cfg

@@ -0,0 +1,7 @@
+[settings]
+line_length = 120
+multi_line_output=3
+include_trailing_comma = true
+ignore_comments = true
+profile = black
+honor_noqa = true

.pre-commit-config.yaml

@@ -0,0 +1,32 @@
+repos:
+
+  - repo: https://github.com/PyCQA/autoflake
+    rev: v2.2.1
+    hooks:
+      - id: autoflake
+        name: autoflake (python)
+        args: ['--in-place', '--remove-unused-variables', '--remove-all-unused-imports', '--ignore-init-module-imports']
+
+  - repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        name: sort all imports (python)
+
+  - repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 23.9.1
+    hooks:
+    - id: black
+      name: black formatter
+      args: ['--line-length=120', '--target-version=py37', '--target-version=py38', '--target-version=py39','--target-version=py310']
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.3.0
+    hooks:
+      - id: check-yaml
+      - id: check-merge-conflict
+      - id: check-case-conflict
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+        args: ['--fix=lf']

README.md

@@ -20,7 +20,7 @@ Here is an example of the captions file:
     },
     {
         "file": "video1.mp4",
-        "captions": ["a  comparison of two opposing team football athletes"]    
+        "captions": ["a  comparison of two opposing team football athletes"]
     }
 ]
 ```
@@ -46,4 +46,4 @@ How to run the script:
 python preprocess_data.py /path/to/captions.json /path/to/video_dir /path/to/output_dir
 ```
 
-Note that this script needs to be run on a machine with a GPU. To avoid CUDA OOM, we filter out the videos that are too long.
+Note that this script needs to be run on a machine with a GPU. To avoid CUDA OOM, we filter out the videos that are too long.

data_utils.py

@@ -1,8 +1,7 @@
 import os
-from typing import Dict, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 from datasets import Dataset as HFDataset
 from datasets import dataset_dict, load_from_disk
@@ -35,9 +34,7 @@ def video2col(video_4d: torch.Tensor, patch_size: int) -> torch.Tensor:
     return torch.stack(out, dim=1).view(-1, c, patch_size, patch_size)
 
 
-def col2video(
-    patches: torch.Tensor, video_shape: Tuple[int, int, int, int]
-) -> torch.Tensor:
+def col2video(patches: torch.Tensor, video_shape: Tuple[int, int, int, int]) -> torch.Tensor:
     """
     Convert a 2D tensor of patches to a 4D video tensor.
 
@@ -74,10 +71,7 @@ def pad_sequences(sequences: List[torch.Tensor]) -> Tuple[torch.Tensor, torch.Te
     """
     max_len = max([sequence.shape[0] for sequence in sequences])
     padded_sequences = [
-        F.pad(
-            sequence, [0] * (sequence.ndim - 1) * 2 + [0, max_len - sequence.shape[0]]
-        )
-        for sequence in sequences
+        F.pad(sequence, [0] * (sequence.ndim - 1) * 2 + [0, max_len - sequence.shape[0]]) for sequence in sequences
     ]
     padded_sequences = torch.stack(padded_sequences, dim=0)
     padding_mask = torch.zeros(
@@ -91,9 +85,7 @@ def pad_sequences(sequences: List[torch.Tensor]) -> Tuple[torch.Tensor, torch.Te
     return padded_sequences, padding_mask
 
 
-def patchify_batch(
-    videos: List[torch.Tensor], patch_size: int
-) -> Tuple[torch.Tensor, torch.Tensor]:
+def patchify_batch(videos: List[torch.Tensor], patch_size: int) -> Tuple[torch.Tensor, torch.Tensor]:
     """Patchify a batch of videos.
 
     Args:
@@ -128,18 +120,14 @@ def make_batch(samples: List[dict], patch_size: int) -> dict:
     }
 
 
-def load_datasets(
-    dataset_paths: Union[PathType, List[PathType]], mode: str = "train"
-) -> Optional[DatasetType]:
+def load_datasets(dataset_paths: Union[PathType, List[PathType]], mode: str = "train") -> Optional[DatasetType]:
     """
     Load pre-tokenized dataset.
     Each instance of dataset is a dictionary with
     `{'input_ids': List[int], 'labels': List[int], sequence: str}` format.
     """
     mode_map = {"train": "train", "dev": "validation", "test": "test"}
-    assert mode in tuple(
-        mode_map
-    ), f"Unsupported mode {mode}, it must be in {tuple(mode_map)}"
+    assert mode in tuple(mode_map), f"Unsupported mode {mode}, it must be in {tuple(mode_map)}"
 
     if isinstance(dataset_paths, (str, os.PathLike)):
         dataset_paths = [dataset_paths]
@@ -148,9 +136,7 @@ def load_datasets(
     for ds_path in dataset_paths:
         ds_path = os.path.abspath(ds_path)
         assert os.path.exists(ds_path), f"Not existed file path {ds_path}"
-        ds_dict = load_from_disk(
-            dataset_path=ds_path, keep_in_memory=False
-        ).with_format("torch")
+        ds_dict = load_from_disk(dataset_path=ds_path, keep_in_memory=False).with_format("torch")
         if isinstance(ds_dict, HFDataset):
             datasets.append(ds_dict)
         else:

diffusion/__init__.py

@@ -9,13 +9,13 @@ from .respace import SpacedDiffusion, space_timesteps
 
 def create_diffusion(
     timestep_respacing,
-    noise_schedule="linear", 
+    noise_schedule="linear",
     use_kl=False,
     sigma_small=False,
     predict_xstart=False,
     learn_sigma=True,
     rescale_learned_sigmas=False,
-    diffusion_steps=1000
+    diffusion_steps=1000,
 ):
     betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps)
     if use_kl:
@@ -29,15 +29,9 @@ def create_diffusion(
     return SpacedDiffusion(
         use_timesteps=space_timesteps(diffusion_steps, timestep_respacing),
         betas=betas,
-        model_mean_type=(
-            gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
-        ),
+        model_mean_type=(gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X),
         model_var_type=(
-            (
-                gd.ModelVarType.FIXED_LARGE
-                if not sigma_small
-                else gd.ModelVarType.FIXED_SMALL
-            )
+            (gd.ModelVarType.FIXED_LARGE if not sigma_small else gd.ModelVarType.FIXED_SMALL)
             if not learn_sigma
             else gd.ModelVarType.LEARNED_RANGE
         ),

diffusion/diffusion_utils.py

@@ -3,8 +3,8 @@
 #     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
 #     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
 
-import torch as th
 import numpy as np
+import torch as th
 
 
 def normal_kl(mean1, logvar1, mean2, logvar2):
@@ -22,18 +22,9 @@ def normal_kl(mean1, logvar1, mean2, logvar2):
 
     # Force variances to be Tensors. Broadcasting helps convert scalars to
     # Tensors, but it does not work for th.exp().
-    logvar1, logvar2 = [
-        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
-        for x in (logvar1, logvar2)
-    ]
+    logvar1, logvar2 = [x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) for x in (logvar1, logvar2)]
 
-    return 0.5 * (
-        -1.0
-        + logvar2
-        - logvar1
-        + th.exp(logvar1 - logvar2)
-        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
-    )
+    return 0.5 * (-1.0 + logvar2 - logvar1 + th.exp(logvar1 - logvar2) + ((mean1 - mean2) ** 2) * th.exp(-logvar2))
 
 
 def approx_standard_normal_cdf(x):

diffusion/gaussian_diffusion.py

@@ -4,11 +4,11 @@
 #     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
 
 
+import enum
 import math
 
 import numpy as np
 import torch as th
-import enum
 
 from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl
 
@@ -45,9 +45,7 @@ class ModelVarType(enum.Enum):
 
 class LossType(enum.Enum):
     MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)
-    RESCALED_MSE = (
-        enum.auto()
-    )  # use raw MSE loss (with RESCALED_KL when learning variances)
+    RESCALED_MSE = enum.auto()  # use raw MSE loss (with RESCALED_KL when learning variances)
     KL = enum.auto()  # use the variational lower-bound
     RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB
 
@@ -70,8 +68,8 @@ def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_time
     if beta_schedule == "quad":
         betas = (
             np.linspace(
-                beta_start ** 0.5,
-                beta_end ** 0.5,
+                beta_start**0.5,
+                beta_end**0.5,
                 num_diffusion_timesteps,
                 dtype=np.float64,
             )
@@ -86,9 +84,7 @@ def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_time
     elif beta_schedule == "const":
         betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
     elif beta_schedule == "jsd":  # 1/T, 1/(T-1), 1/(T-2), ..., 1
-        betas = 1.0 / np.linspace(
-            num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64
-        )
+        betas = 1.0 / np.linspace(num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64)
     else:
         raise NotImplementedError(beta_schedule)
     assert betas.shape == (num_diffusion_timesteps,)
@@ -150,15 +146,7 @@ class GaussianDiffusion:
                   starting at T and going to 1.
     """
 
-    def __init__(
-        self,
-        *,
-        betas,
-        model_mean_type,
-        model_var_type,
-        loss_type
-    ):
-
+    def __init__(self, *, betas, model_mean_type, model_var_type, loss_type):
         self.model_mean_type = model_mean_type
         self.model_var_type = model_var_type
         self.loss_type = loss_type
@@ -185,20 +173,16 @@ class GaussianDiffusion:
         self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
 
         # calculations for posterior q(x_{t-1} | x_t, x_0)
-        self.posterior_variance = (
-            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
-        )
+        self.posterior_variance = betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
         # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
-        self.posterior_log_variance_clipped = np.log(
-            np.append(self.posterior_variance[1], self.posterior_variance[1:])
-        ) if len(self.posterior_variance) > 1 else np.array([])
+        self.posterior_log_variance_clipped = (
+            np.log(np.append(self.posterior_variance[1], self.posterior_variance[1:]))
+            if len(self.posterior_variance) > 1
+            else np.array([])
+        )
 
-        self.posterior_mean_coef1 = (
-            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
-        )
-        self.posterior_mean_coef2 = (
-            (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)
-        )
+        self.posterior_mean_coef1 = betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        self.posterior_mean_coef2 = (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)
 
     def q_mean_variance(self, x_start, t):
         """
@@ -240,9 +224,7 @@ class GaussianDiffusion:
             + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
         )
         posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
-        posterior_log_variance_clipped = _extract_into_tensor(
-            self.posterior_log_variance_clipped, t, x_t.shape
-        )
+        posterior_log_variance_clipped = _extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
         assert (
             posterior_mean.shape[0]
             == posterior_variance.shape[0]
@@ -317,9 +299,7 @@ class GaussianDiffusion:
         if self.model_mean_type == ModelMeanType.START_X:
             pred_xstart = process_xstart(model_output)
         else:
-            pred_xstart = process_xstart(
-                self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
-            )
+            pred_xstart = process_xstart(self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output))
         model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)
 
         assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
@@ -408,9 +388,7 @@ class GaussianDiffusion:
             model_kwargs=model_kwargs,
         )
         noise = th.randn_like(x)
-        nonzero_mask = (
-            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
-        )  # no noise when t == 0
+        nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))  # no noise when t == 0
         if cond_fn is not None:
             out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs)
         sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
@@ -542,20 +520,11 @@ class GaussianDiffusion:
 
         alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
         alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
-        sigma = (
-            eta
-            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
-            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
-        )
+        sigma = eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * th.sqrt(1 - alpha_bar / alpha_bar_prev)
         # Equation 12.
         noise = th.randn_like(x)
-        mean_pred = (
-            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
-            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
-        )
-        nonzero_mask = (
-            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
-        )  # no noise when t == 0
+        mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_prev) + th.sqrt(1 - alpha_bar_prev - sigma**2) * eps
+        nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))  # no noise when t == 0
         sample = mean_pred + nonzero_mask * sigma * noise
         return {"sample": sample, "pred_xstart": out["pred_xstart"]}
 
@@ -587,8 +556,7 @@ class GaussianDiffusion:
         # Usually our model outputs epsilon, but we re-derive it
         # in case we used x_start or x_prev prediction.
         eps = (
-            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
-            - out["pred_xstart"]
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x - out["pred_xstart"]
         ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
         alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
 
@@ -679,9 +647,7 @@ class GaussianDiffusion:
                 yield out
                 img = out["sample"]
 
-    def _vb_terms_bpd(
-            self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
-    ):
+    def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None):
         """
         Get a term for the variational lower-bound.
         The resulting units are bits (rather than nats, as one might expect).
@@ -690,15 +656,9 @@ class GaussianDiffusion:
                  - 'output': a shape [N] tensor of NLLs or KLs.
                  - 'pred_xstart': the x_0 predictions.
         """
-        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
-            x_start=x_start, x_t=x_t, t=t
-        )
-        out = self.p_mean_variance(
-            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
-        )
-        kl = normal_kl(
-            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
-        )
+        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)
+        out = self.p_mean_variance(model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs)
+        kl = normal_kl(true_mean, true_log_variance_clipped, out["mean"], out["log_variance"])
         kl = mean_flat(kl) / np.log(2.0)
 
         decoder_nll = -discretized_gaussian_log_likelihood(
@@ -769,9 +729,7 @@ class GaussianDiffusion:
                     terms["vb"] *= self.num_timesteps / 1000.0
 
             target = {
-                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
-                    x_start=x_start, x_t=x_t, t=t
-                )[0],
+                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)[0],
                 ModelMeanType.START_X: x_start,
                 ModelMeanType.EPSILON: noise,
             }[self.model_mean_type]
@@ -797,9 +755,7 @@ class GaussianDiffusion:
         batch_size = x_start.shape[0]
         t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
         qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
-        kl_prior = normal_kl(
-            mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
-        )
+        kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
         return mean_flat(kl_prior) / np.log(2.0)
 
     def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):

diffusion/respace.py

@@ -34,9 +34,7 @@ def space_timesteps(num_timesteps, section_counts):
             for i in range(1, num_timesteps):
                 if len(range(0, num_timesteps, i)) == desired_count:
                     return set(range(0, num_timesteps, i))
-            raise ValueError(
-                f"cannot create exactly {num_timesteps} steps with an integer stride"
-            )
+            raise ValueError(f"cannot create exactly {num_timesteps} steps with an integer stride")
         section_counts = [int(x) for x in section_counts.split(",")]
     size_per = num_timesteps // len(section_counts)
     extra = num_timesteps % len(section_counts)
@@ -45,9 +43,7 @@ def space_timesteps(num_timesteps, section_counts):
     for i, section_count in enumerate(section_counts):
         size = size_per + (1 if i < extra else 0)
         if size < section_count:
-            raise ValueError(
-                f"cannot divide section of {size} steps into {section_count}"
-            )
+            raise ValueError(f"cannot divide section of {size} steps into {section_count}")
         if section_count <= 1:
             frac_stride = 1
         else:
@@ -86,14 +82,10 @@ class SpacedDiffusion(GaussianDiffusion):
         kwargs["betas"] = np.array(new_betas)
         super().__init__(**kwargs)
 
-    def p_mean_variance(
-        self, model, *args, **kwargs
-    ):  # pylint: disable=signature-differs
+    def p_mean_variance(self, model, *args, **kwargs):  # pylint: disable=signature-differs
         return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
 
-    def training_losses(
-        self, model, *args, **kwargs
-    ):  # pylint: disable=signature-differs
+    def training_losses(self, model, *args, **kwargs):  # pylint: disable=signature-differs
         return super().training_losses(self._wrap_model(model), *args, **kwargs)
 
     def condition_mean(self, cond_fn, *args, **kwargs):
@@ -105,9 +97,7 @@ class SpacedDiffusion(GaussianDiffusion):
     def _wrap_model(self, model):
         if isinstance(model, _WrappedModel):
             return model
-        return _WrappedModel(
-            model, self.timestep_map, self.original_num_steps
-        )
+        return _WrappedModel(model, self.timestep_map, self.original_num_steps)
 
     def _scale_timesteps(self, t):
         # Scaling is done by the wrapped model.

diffusion/timestep_sampler.py

@@ -79,10 +79,7 @@ class LossAwareSampler(ScheduleSampler):
         :param local_ts: an integer Tensor of timesteps.
         :param local_losses: a 1D Tensor of losses.
         """
-        batch_sizes = [
-            th.tensor([0], dtype=th.int32, device=local_ts.device)
-            for _ in range(dist.get_world_size())
-        ]
+        batch_sizes = [th.tensor([0], dtype=th.int32, device=local_ts.device) for _ in range(dist.get_world_size())]
         dist.all_gather(
             batch_sizes,
             th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device),
@@ -96,9 +93,7 @@ class LossAwareSampler(ScheduleSampler):
         loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes]
         dist.all_gather(timestep_batches, local_ts)
         dist.all_gather(loss_batches, local_losses)
-        timesteps = [
-            x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]
-        ]
+        timesteps = [x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]]
         losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]]
         self.update_with_all_losses(timesteps, losses)
 
@@ -122,15 +117,13 @@ class LossSecondMomentResampler(LossAwareSampler):
         self.diffusion = diffusion
         self.history_per_term = history_per_term
         self.uniform_prob = uniform_prob
-        self._loss_history = np.zeros(
-            [diffusion.num_timesteps, history_per_term], dtype=np.float64
-        )
+        self._loss_history = np.zeros([diffusion.num_timesteps, history_per_term], dtype=np.float64)
         self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
 
     def weights(self):
         if not self._warmed_up():
             return np.ones([self.diffusion.num_timesteps], dtype=np.float64)
-        weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))
+        weights = np.sqrt(np.mean(self._loss_history**2, axis=-1))
         weights /= np.sum(weights)
         weights *= 1 - self.uniform_prob
         weights += self.uniform_prob / len(weights)

download.py

@@ -7,12 +7,12 @@
 """
 Functions for downloading pre-trained DiT models
 """
-from torchvision.datasets.utils import download_url
-import torch
 import os
 
+import torch
+from torchvision.datasets.utils import download_url
 
-pretrained_models = {'DiT-XL-2-512x512.pt', 'DiT-XL-2-256x256.pt'}
+pretrained_models = {"DiT-XL-2-512x512.pt", "DiT-XL-2-256x256.pt"}
 
 
 def find_model(model_name):
@@ -22,7 +22,7 @@ def find_model(model_name):
     if model_name in pretrained_models:  # Find/download our pre-trained DiT checkpoints
         return download_model(model_name)
     else:  # Load a custom DiT checkpoint:
-        assert os.path.isfile(model_name), f'Could not find DiT checkpoint at {model_name}'
+        assert os.path.isfile(model_name), f"Could not find DiT checkpoint at {model_name}"
         checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
         if "ema" in checkpoint:  # supports checkpoints from train.py
             checkpoint = checkpoint["ema"]
@@ -34,11 +34,11 @@ def download_model(model_name):
     Downloads a pre-trained DiT model from the web.
     """
     assert model_name in pretrained_models
-    local_path = f'pretrained_models/{model_name}'
+    local_path = f"pretrained_models/{model_name}"
     if not os.path.isfile(local_path):
-        os.makedirs('pretrained_models', exist_ok=True)
-        web_path = f'https://dl.fbaipublicfiles.com/DiT/models/{model_name}'
-        download_url(web_path, 'pretrained_models')
+        os.makedirs("pretrained_models", exist_ok=True)
+        web_path = f"https://dl.fbaipublicfiles.com/DiT/models/{model_name}"
+        download_url(web_path, "pretrained_models")
     model = torch.load(local_path, map_location=lambda storage, loc: storage)
     return model
 
@@ -47,4 +47,4 @@ if __name__ == "__main__":
     # Download all DiT checkpoints
     for model in pretrained_models:
         download_model(model)
-    print('Done.')
+    print("Done.")

models.py

@@ -16,7 +16,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from timm.models.vision_transformer import Attention, Mlp, PatchEmbed
+from timm.models.vision_transformer import Attention, Mlp
 
 
 def modulate(x, shift, scale):
@@ -54,17 +54,13 @@ class TimestepEmbedder(nn.Module):
         """
         # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
         half = dim // 2
-        freqs = torch.exp(
-            -math.log(max_period)
-            * torch.arange(start=0, end=half, dtype=torch.float32)
-            / half
-        ).to(device=t.device)
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
+            device=t.device
+        )
         args = t[:, None].float() * freqs[None]
         embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
         if dim % 2:
-            embedding = torch.cat(
-                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
-            )
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
         return embedding
 
     def forward(self, t):
@@ -81,9 +77,7 @@ class LabelEmbedder(nn.Module):
     def __init__(self, num_classes, hidden_size, dropout_prob):
         super().__init__()
         use_cfg_embedding = dropout_prob > 0
-        self.embedding_table = nn.Embedding(
-            num_classes + use_cfg_embedding, hidden_size
-        )
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
         self.num_classes = num_classes
         self.dropout_prob = dropout_prob
 
@@ -92,9 +86,7 @@ class LabelEmbedder(nn.Module):
         Drops labels to enable classifier-free guidance.
         """
         if force_drop_ids is None:
-            drop_ids = (
-                torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
-            )
+            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
         else:
             drop_ids = force_drop_ids == 1
         labels = torch.where(drop_ids, self.num_classes, labels)
@@ -121,26 +113,20 @@ class PatchEmbedder(nn.Module):
     ) -> None:
         super().__init__()
         self.patch_size = patch_size
-        self.proj = nn.Conv2d(
-            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias
-        )
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
         self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # [B, S, C, P, P] -> [B, S, C*P*P]
         x = x.view(*x.shape[:2], -1)
-        out = F.linear(
-            x, self.proj.weight.view(self.proj.weight.shape[0], -1), self.proj.bias
-        )
+        out = F.linear(x, self.proj.weight.view(self.proj.weight.shape[0], -1), self.proj.bias)
         out = self.norm(out)
         # [B, S, H]
         return out
 
 
 class TextEmbedder(nn.Module):
-    def __init__(
-        self, in_features: int, embed_dim: int = 768, bias: bool = True
-    ) -> None:
+    def __init__(self, in_features: int, embed_dim: int = 768, bias: bool = True) -> None:
         super().__init__()
         self.proj = nn.Linear(in_features, embed_dim, bias=bias)
 
@@ -158,9 +144,7 @@ class PositionEmbedding(nn.Module):
 
     def _set_pos_embed_cache(self, seq_len: int):
         self.max_seq_len_cached = seq_len
-        pos_embed = get_2d_sincos_pos_embed(
-            self.dim, int(self.max_position_embeddings**0.5)
-        )
+        pos_embed = get_2d_sincos_pos_embed(self.dim, int(self.max_position_embeddings**0.5))
         pos_embed = torch.from_numpy(pos_embed).float()
         # [S, H]
         self.register_buffer("pos_embed_cache", pos_embed, persistent=False)
@@ -187,9 +171,7 @@ class DiTBlock(nn.Module):
     def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
         super().__init__()
         self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.attn = Attention(
-            hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs
-        )
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
         self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
         mlp_hidden_dim = int(hidden_size * mlp_ratio)
         approx_gelu = lambda: nn.GELU(approximate="tanh")
@@ -214,9 +196,7 @@ class FinalLayer(nn.Module):
 
     def __init__(self, hidden_size, patch_size, out_channels):
         super().__init__()
-        self.linear = nn.Linear(
-            hidden_size, patch_size * patch_size * out_channels, bias=True
-        )
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
         self.patch_size = patch_size
 
     def unpatchify(self, x):
@@ -254,19 +234,12 @@ class DiT(nn.Module):
         self.patch_size = patch_size
         self.num_heads = num_heads
 
-        self.video_embedder = PatchEmbedder(
-            patch_size, in_channels, hidden_size, bias=True
-        )
+        self.video_embedder = PatchEmbedder(patch_size, in_channels, hidden_size, bias=True)
         self.t_embedder = TimestepEmbedder(hidden_size)
         self.text_embedder = TextEmbedder(text_embed_dim, hidden_size, bias=True)
         self.pos_embed = PositionEmbedding(hidden_size, max_num_embeddings)
 
-        self.blocks = nn.ModuleList(
-            [
-                DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio)
-                for _ in range(depth)
-            ]
-        )
+        self.blocks = nn.ModuleList([DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth)])
         self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
         self.initialize_weights()
 
@@ -353,9 +326,7 @@ def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=
     grid = grid.reshape([2, 1, grid_size, grid_size])
     pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
     if cls_token and extra_tokens > 0:
-        pos_embed = np.concatenate(
-            [np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0
-        )
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
     return pos_embed
 
 

preprocess_data.py

@@ -9,6 +9,7 @@ from transformers import AutoModel, AutoTokenizer, CLIPTextModel
 
 EMPTY_SAMPLE = {"video_file": [], "video_latent_states": [], "text_latent_states": []}
 
+
 def preprocess_video(video):
     # [T, H, W, C] to [C, T, H, W]
     video = video.permute(3, 0, 1, 2)
@@ -17,6 +18,7 @@ def preprocess_video(video):
     video = video / 255 - 0.5
     return video.unsqueeze(0)
 
+
 def process_video(video_path, vqvae):
     video = read_video(video_path, pts_unit="sec")[0]
     video = preprocess_video(video)
@@ -25,6 +27,7 @@ def process_video(video_path, vqvae):
     latent_states = vqvae.encode(video)
     return latent_states.squeeze(0).tolist()
 
+
 def process_text(text, tokenizer, text_model):
     inputs = tokenizer(text, padding=True, return_tensors="pt")
     inputs = {k: v.cuda() for k, v in inputs.items()}
@@ -35,12 +38,13 @@ def process_text(text, tokenizer, text_model):
         output_states.append(valid_x.tolist())
     return output_states
 
+
 @torch.no_grad()
 def process_item(item, video_dir, tokenizer, text_model, vqvae):
     video_path = os.path.join(video_dir, item["file"])
     try:
         video_latent_states = process_video(video_path, vqvae)
-    except ValueError as e:
+    except ValueError:
         return EMPTY_SAMPLE
     torch.cuda.empty_cache()
     text_latent_states = process_text(item["captions"], tokenizer, text_model)
@@ -48,14 +52,23 @@ def process_item(item, video_dir, tokenizer, text_model, vqvae):
     return {
         "video_file": [item["file"]] * len(text_latent_states),
         "video_latent_states": [video_latent_states] * len(text_latent_states),
-        "text_latent_states": text_latent_states
+        "text_latent_states": text_latent_states,
     }
 
+
 def process_batch(batch, video_dir, tokenizer, text_model, vqvae):
     item = {"file": batch["file"][0], "captions": batch["captions"][0]}
     return process_item(item, video_dir, tokenizer, text_model, vqvae)
 
-def process_dataset(captions_file, video_dir, output_dir, num_spliced_dataset_bins=10, text_model="openai/clip-vit-base-patch32", vae_model="hpcai-tech/vqvae"):
+
+def process_dataset(
+    captions_file,
+    video_dir,
+    output_dir,
+    num_spliced_dataset_bins=10,
+    text_model="openai/clip-vit-base-patch32",
+    vae_model="hpcai-tech/vqvae",
+):
     tokenizer = AutoTokenizer.from_pretrained(text_model)
     text_model = CLIPTextModel.from_pretrained(text_model).cuda().eval()
     vqvae = AutoModel.from_pretrained(vae_model, trust_remote_code=True).cuda().eval()
@@ -72,33 +85,41 @@ def process_dataset(captions_file, video_dir, output_dir, num_spliced_dataset_bi
         if end > 100:
             end = 100
         train_splits.append(f"train[{start}%:{end}%]")
-    
+
     ds = load_dataset("json", data_files=captions_file, keep_in_memory=False, split=train_splits)
 
     for i, part_ds in enumerate(ds):
         print(f"Processing part {i+1}/{len(ds)}")
-        part_ds = part_ds.map(process_batch, 
-                              fn_kwargs={
-                                  "video_dir": video_dir, 
-                                  "tokenizer": tokenizer, 
-                                  "text_model": text_model, 
-                                  "vqvae": vqvae
-                                  }, 
-                                  batched=True, 
-                                  batch_size=1,
-                                  keep_in_memory=False,
-            remove_columns=part_ds.column_names)
+        part_ds = part_ds.map(
+            process_batch,
+            fn_kwargs={"video_dir": video_dir, "tokenizer": tokenizer, "text_model": text_model, "vqvae": vqvae},
+            batched=True,
+            batch_size=1,
+            keep_in_memory=False,
+            remove_columns=part_ds.column_names,
+        )
         output_path = os.path.join(output_dir, f"part-{i:05d}")
         part_ds.save_to_disk(output_path)
 
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Preprocess data')
-    parser.add_argument("captions_file", type=str, help="Path to the captions file. It should be a JSON file or a JSONL file")
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Preprocess data")
+    parser.add_argument(
+        "captions_file", type=str, help="Path to the captions file. It should be a JSON file or a JSONL file"
+    )
     parser.add_argument("video_dir", type=str, help="Path to the video directory")
     parser.add_argument("output_dir", type=str, help="Path to the output directory")
-    parser.add_argument("-n", "--num_spliced_dataset_bins", type=int, default=10, help="Number of bins for spliced dataset")
+    parser.add_argument(
+        "-n", "--num_spliced_dataset_bins", type=int, default=10, help="Number of bins for spliced dataset"
+    )
     parser.add_argument("--text_model", type=str, default="openai/clip-vit-base-patch32", help="CLIP text model")
     parser.add_argument("--vae_model", type=str, default="hpcai-tech/vqvae", help="VQ-VAE model")
     args = parser.parse_args()
-    process_dataset(args.captions_file, args.video_dir, args.output_dir, args.num_spliced_dataset_bins, args.text_model, args.vae_model)
+    process_dataset(
+        args.captions_file,
+        args.video_dir,
+        args.output_dir,
+        args.num_spliced_dataset_bins,
+        args.text_model,
+        args.vae_model,
+    )

sample.py

@@ -8,14 +8,17 @@
 Sample new images from a pre-trained DiT.
 """
 import torch
+
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
-from torchvision.utils import save_image
-from diffusion import create_diffusion
+import argparse
+
 from diffusers.models import AutoencoderKL
+from torchvision.utils import save_image
+
+from diffusion import create_diffusion
 from download import find_model
 from models import DiT_models
-import argparse
 
 
 def main(args):
@@ -31,10 +34,7 @@ def main(args):
 
     # Load model:
     latent_size = args.image_size // 8
-    model = DiT_models[args.model](
-        input_size=latent_size,
-        num_classes=args.num_classes
-    ).to(device)
+    model = DiT_models[args.model](input_size=latent_size, num_classes=args.num_classes).to(device)
     # Auto-download a pre-trained model or load a custom DiT checkpoint from train.py:
     ckpt_path = args.ckpt or f"DiT-XL-2-{args.image_size}x{args.image_size}.pt"
     state_dict = find_model(ckpt_path)
@@ -77,7 +77,11 @@ if __name__ == "__main__":
     parser.add_argument("--cfg-scale", type=float, default=4.0)
     parser.add_argument("--num-sampling-steps", type=int, default=250)
     parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--ckpt", type=str, default=None,
-                        help="Optional path to a DiT checkpoint (default: auto-download a pre-trained DiT-XL/2 model).")
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        default=None,
+        help="Optional path to a DiT checkpoint (default: auto-download a pre-trained DiT-XL/2 model).",
+    )
     args = parser.parse_args()
     main(args)

train.py

@@ -25,7 +25,6 @@ from colossalai.booster.plugin import LowLevelZeroPlugin
 from colossalai.cluster import DistCoordinator
 from colossalai.logging import get_dist_logger
 from colossalai.utils import get_current_device
-from diffusers.models import AutoencoderKL
 from tqdm import tqdm
 
 from data_utils import load_datasets, make_batch
@@ -99,9 +98,7 @@ def main(args):
     model.train()  # important! This enables embedding dropout for classifier-free guidance
     ema.eval()  # EMA model should always be in eval mode
 
-    diffusion = create_diffusion(
-        timestep_respacing=""
-    )  # default: 1000 steps, linear noise schedule
+    diffusion = create_diffusion(timestep_respacing="")  # default: 1000 steps, linear noise schedule
 
     # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
     opt = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0)
@@ -153,19 +150,11 @@ def main(args):
                 pbar.update()
 
                 # Save DiT checkpoint:
-                if (
-                    args.save_interval > 0 and (step + 1) % args.save_interval == 0
-                ) or (step + 1) == len(dataloader):
-                    save_path = os.path.join(
-                        args.checkpoint_dir, f"epoch-{epoch}-step-{step}"
-                    )
+                if (args.save_interval > 0 and (step + 1) % args.save_interval == 0) or (step + 1) == len(dataloader):
+                    save_path = os.path.join(args.checkpoint_dir, f"epoch-{epoch}-step-{step}")
                     os.makedirs(save_path, exist_ok=True)
-                    booster.save_model(
-                        model, os.path.join(save_path, "model"), shard=True
-                    )
-                    booster.save_optimizer(
-                        opt, os.path.join(save_path, "optimizer"), shard=True
-                    )
+                    booster.save_model(model, os.path.join(save_path, "model"), shard=True)
+                    booster.save_optimizer(opt, os.path.join(save_path, "optimizer"), shard=True)
                     if coordinator.is_master():
                         ema_state_dict = ema.state_dict()
                         for k, v in ema_state_dict.items():
@@ -179,9 +168,7 @@ def main(args):
 if __name__ == "__main__":
     # Default args here will train DiT-XL/2 with the hyperparameters we used in our paper (except training iters).
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-m", "--model", type=str, choices=list(DiT_models.keys()), default="DiT-S/8"
-    )
+    parser.add_argument("-m", "--model", type=str, choices=list(DiT_models.keys()), default="DiT-S/8")
     parser.add_argument("--dataset", nargs="+", default=[])
     parser.add_argument("-e", "--epochs", type=int, default=10)
     parser.add_argument("-b", "--batch_size", type=int, default=4)