From 78dfc3064fda17ae031809fe8833090128613588 Mon Sep 17 00:00:00 2001 From: "snorkell-ai[bot]" <146478655+snorkell-ai[bot]@users.noreply.github.com> Date: Tue, 27 Feb 2024 10:07:56 +0000 Subject: [PATCH] Generated Documentation --- diffusion/__init__.py | 18 ++ diffusion/diffusion_utils.py | 61 ++-- diffusion/gaussian_diffusion.py | 521 +++++++++++++++++++++++--------- diffusion/timestep_sampler.py | 119 ++++++-- download.py | 25 +- models.py | 365 ++++++++++++++++++++-- sample.py | 12 + sample_ddp.py | 25 +- train.py | 62 +++- 9 files changed, 974 insertions(+), 234 deletions(-) diff --git a/diffusion/__init__.py b/diffusion/__init__.py index 8c536a98..6acea824 100644 --- a/diffusion/__init__.py +++ b/diffusion/__init__.py @@ -17,6 +17,24 @@ def create_diffusion( rescale_learned_sigmas=False, diffusion_steps=1000 ): + """ Create a diffusion model for the given parameters. + + This function creates a diffusion model based on the provided parameters, including the timestep respacing, noise schedule, use of KL divergence, sigma size, prediction of xstart, learning of sigma, rescaling of learned sigmas, and diffusion steps. + + Args: + timestep_respacing (int or list): The respacing of timesteps for diffusion. + noise_schedule (str?): The schedule for noise. Defaults to "linear". + use_kl (bool?): Whether to use KL divergence. Defaults to False. + sigma_small (bool?): Whether sigma is small. Defaults to False. + predict_xstart (bool?): Whether to predict xstart. Defaults to False. + learn_sigma (bool?): Whether to learn sigma. Defaults to True. + rescale_learned_sigmas (bool?): Whether to rescale learned sigmas. Defaults to False. + diffusion_steps (int?): The number of diffusion steps. Defaults to 1000. + + Returns: + SpacedDiffusion: A diffusion model based on the provided parameters. + """ + betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps) if use_kl: loss_type = gd.LossType.RESCALED_KL diff --git a/diffusion/diffusion_utils.py b/diffusion/diffusion_utils.py index e493a6a3..4a582db0 100644 --- a/diffusion/diffusion_utils.py +++ b/diffusion/diffusion_utils.py @@ -8,10 +8,21 @@ def normal_kl(mean1, logvar1, mean2, logvar2): - """ - Compute the KL divergence between two gaussians. - Shapes are automatically broadcasted, so batches can be compared to - scalars, among other use cases. + """ Compute the KL divergence between two Gaussian distributions. + + This function computes the Kullback-Leibler (KL) divergence between two Gaussian distributions. The shapes of the input parameters are automatically broadcasted, allowing for comparisons between batches and scalars. + + Args: + mean1 (Tensor or float): The mean of the first Gaussian distribution. + logvar1 (Tensor or float): The log variance of the first Gaussian distribution. + mean2 (Tensor or float): The mean of the second Gaussian distribution. + logvar2 (Tensor or float): The log variance of the second Gaussian distribution. + + Returns: + Tensor: The computed KL divergence between the two Gaussian distributions. + + Raises: + AssertionError: If all input arguments are not Tensors. """ tensor = None for obj in (mean1, logvar1, mean2, logvar2): @@ -37,20 +48,27 @@ def normal_kl(mean1, logvar1, mean2, logvar2): def approx_standard_normal_cdf(x): - """ - A fast approximation of the cumulative distribution function of the - standard normal. + """ A fast approximation of the cumulative distribution function of the standard normal. + + Args: + x (float): The input value for the standard normal distribution. + + Returns: + float: The approximate cumulative distribution function value. """ return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) def continuous_gaussian_log_likelihood(x, *, means, log_scales): - """ - Compute the log-likelihood of a continuous Gaussian distribution. - :param x: the targets - :param means: the Gaussian mean Tensor. - :param log_scales: the Gaussian log stddev Tensor. - :return: a tensor like x of log probabilities (in nats). + """ Compute the log-likelihood of a continuous Gaussian distribution. + + Args: + x (tensor): The targets. + means (tensor): The Gaussian mean Tensor. + log_scales (tensor): The Gaussian log stddev Tensor. + + Returns: + tensor: A tensor like x of log probabilities (in nats). """ centered_x = x - means inv_stdv = th.exp(-log_scales) @@ -60,14 +78,17 @@ def continuous_gaussian_log_likelihood(x, *, means, log_scales): def discretized_gaussian_log_likelihood(x, *, means, log_scales): - """ - Compute the log-likelihood of a Gaussian distribution discretizing to a + """ Compute the log-likelihood of a Gaussian distribution discretizing to a given image. - :param x: the target images. It is assumed that this was uint8 values, - rescaled to the range [-1, 1]. - :param means: the Gaussian mean Tensor. - :param log_scales: the Gaussian log stddev Tensor. - :return: a tensor like x of log probabilities (in nats). + + Args: + x (Tensor): The target images. It is assumed that this was uint8 values, + rescaled to the range [-1, 1]. + means (Tensor): The Gaussian mean Tensor. + log_scales (Tensor): The Gaussian log stddev Tensor. + + Returns: + Tensor: A tensor like x of log probabilities (in nats). """ assert x.shape == means.shape == log_scales.shape centered_x = x - means diff --git a/diffusion/gaussian_diffusion.py b/diffusion/gaussian_diffusion.py index ccbcefec..b5153f2c 100644 --- a/diffusion/gaussian_diffusion.py +++ b/diffusion/gaussian_diffusion.py @@ -14,8 +14,13 @@ def mean_flat(tensor): - """ - Take the mean over all non-batch dimensions. + """ Take the mean over all non-batch dimensions. + + Args: + tensor (tensor): A tensor containing numeric values. + + Returns: + tensor: The mean of the input tensor over non-batch dimensions. """ return tensor.mean(dim=list(range(1, len(tensor.shape)))) @@ -52,10 +57,31 @@ class LossType(enum.Enum): RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB def is_vb(self): + """ Check if the LossType is either KL or RESCALED_KL. + + Returns: + bool: True if the LossType is KL or RESCALED_KL, False otherwise. + """ + return self == LossType.KL or self == LossType.RESCALED_KL def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac): + """ Warm up the beta values for diffusion process. + + This function initializes the beta values for the diffusion process with a warm-up period, where the beta values + transition from `beta_start` to `beta_end` over a specified fraction of the total diffusion timesteps. + + Args: + beta_start (float): The initial value of beta. + beta_end (float): The final value of beta. + num_diffusion_timesteps (int): The total number of diffusion timesteps. + warmup_frac (float): The fraction of diffusion timesteps over which the warm-up occurs. + + Returns: + numpy.ndarray: An array of beta values for each diffusion timestep. + """ + betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) warmup_time = int(num_diffusion_timesteps * warmup_frac) betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64) @@ -63,9 +89,20 @@ def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac): def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps): - """ - This is the deprecated API for creating beta schedules. + """ This is the deprecated API for creating beta schedules. See get_named_beta_schedule() for the new library of schedules. + + Args: + beta_schedule (str): The type of beta schedule to be used. + beta_start (float): The starting value of beta. + beta_end (float): The ending value of beta. + num_diffusion_timesteps (int): The number of diffusion timesteps. + + Returns: + numpy.ndarray: An array of beta values for the specified schedule. + + Raises: + NotImplementedError: If the specified beta schedule is not implemented. """ if beta_schedule == "quad": betas = ( @@ -96,12 +133,22 @@ def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_time def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): - """ - Get a pre-defined beta schedule for the given name. + """ Get a pre-defined beta schedule for the given name. + The beta schedule library consists of beta schedules which remain similar - in the limit of num_diffusion_timesteps. - Beta schedules may be added, but should not be removed or changed once - they are committed to maintain backwards compatibility. + in the limit of num_diffusion_timesteps. Beta schedules may be added, but + should not be removed or changed once they are committed to maintain + backwards compatibility. + + Args: + schedule_name (str): The name of the beta schedule. + num_diffusion_timesteps (int): The number of diffusion timesteps. + + Returns: + list: A pre-defined beta schedule for the given name. + + Raises: + NotImplementedError: If the input schedule_name is not recognized. """ if schedule_name == "linear": # Linear schedule from Ho et al, extended to work for any number of @@ -123,15 +170,19 @@ def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): - """ - Create a beta schedule that discretizes the given alpha_t_bar function, + """ Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of (1-beta) over time from t = [0,1]. - :param num_diffusion_timesteps: the number of betas to produce. - :param alpha_bar: a lambda that takes an argument t from 0 to 1 and - produces the cumulative product of (1-beta) up to that - part of the diffusion process. - :param max_beta: the maximum beta to use; use values lower than 1 to - prevent singularities. + + Args: + num_diffusion_timesteps (int): The number of betas to produce. + alpha_bar (function): A lambda that takes an argument t from 0 to 1 and + produces the cumulative product of (1-beta) up to that + part of the diffusion process. + max_beta (float?): The maximum beta to use; use values lower than 1 to + prevent singularities. + + Returns: + numpy.ndarray: An array of beta values. """ betas = [] for i in range(num_diffusion_timesteps): @@ -158,6 +209,21 @@ def __init__( model_var_type, loss_type ): + """ Initializes the object with the specified parameters. + + This method initializes the object with the provided parameters and performs necessary calculations for further use. + + Args: + betas (list): A list of beta values. + model_mean_type (str): The type of model mean. + model_var_type (str): The type of model variance. + loss_type (str): The type of loss. + + + Raises: + AssertionError: If the betas are not 1-D, or if any beta value is not within the range (0, 1). + """ + self.model_mean_type = model_mean_type self.model_var_type = model_var_type @@ -201,11 +267,14 @@ def __init__( ) def q_mean_variance(self, x_start, t): - """ - Get the distribution q(x_t | x_0). - :param x_start: the [N x C x ...] tensor of noiseless inputs. - :param t: the number of diffusion steps (minus 1). Here, 0 means one step. - :return: A tuple (mean, variance, log_variance), all of x_start's shape. + """ Get the distribution q(x_t | x_0). + + Args: + x_start (tensor): The [N x C x ...] tensor of noiseless inputs. + t (int): The number of diffusion steps (minus 1). Here, 0 means one step. + + Returns: + tuple: A tuple (mean, variance, log_variance), all of x_start's shape. """ mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) @@ -213,13 +282,16 @@ def q_mean_variance(self, x_start, t): return mean, variance, log_variance def q_sample(self, x_start, t, noise=None): - """ - Diffuse the data for a given number of diffusion steps. + """ Diffuse the data for a given number of diffusion steps. In other words, sample from q(x_t | x_0). - :param x_start: the initial data batch. - :param t: the number of diffusion steps (minus 1). Here, 0 means one step. - :param noise: if specified, the split-out normal noise. - :return: A noisy version of x_start. + + Args: + x_start (tensor): The initial data batch. + t (int): The number of diffusion steps (minus 1). Here, 0 means one step. + noise (tensor?): If specified, the split-out normal noise. + + Returns: + tensor: A noisy version of x_start. """ if noise is None: noise = th.randn_like(x_start) @@ -230,9 +302,19 @@ def q_sample(self, x_start, t, noise=None): ) def q_posterior_mean_variance(self, x_start, x_t, t): - """ - Compute the mean and variance of the diffusion posterior: - q(x_{t-1} | x_t, x_0) + """ Compute the mean and variance of the diffusion posterior: + q(x_{t-1} | x_t, x_0) + + Args: + x_start (tensor): The initial state. + x_t (tensor): The state at time t. + t (int): The time step. + + Returns: + tuple: A tuple containing: + - posterior_mean (tensor): The mean of the diffusion posterior. + - posterior_variance (tensor): The variance of the diffusion posterior. + - posterior_log_variance_clipped (tensor): The clipped log variance of the diffusion posterior. """ assert x_start.shape == x_t.shape posterior_mean = ( @@ -252,24 +334,28 @@ def q_posterior_mean_variance(self, x_start, x_t, t): return posterior_mean, posterior_variance, posterior_log_variance_clipped def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None): - """ - Apply the model to get p(x_{t-1} | x_t), as well as a prediction of + """ Apply the model to get p(x_{t-1} | x_t), as well as a prediction of the initial x, x_0. - :param model: the model, which takes a signal and a batch of timesteps - as input. - :param x: the [N x C x ...] tensor at time t. - :param t: a 1-D Tensor of timesteps. - :param clip_denoised: if True, clip the denoised signal into [-1, 1]. - :param denoised_fn: if not None, a function which applies to the - x_start prediction before it is used to sample. Applies before - clip_denoised. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :return: a dict with the following keys: - - 'mean': the model mean output. - - 'variance': the model variance output. - - 'log_variance': the log of 'variance'. - - 'pred_xstart': the prediction for x_0. + + Args: + model: object - the model, which takes a signal and a batch of timesteps as input. + x: tensor - the [N x C x ...] tensor at time t. + t: tensor - a 1-D Tensor of timesteps. + clip_denoised: bool - if True, clip the denoised signal into [-1, 1]. + denoised_fn: function - if not None, a function which applies to the x_start prediction before it is used to sample. Applies before clip_denoised. + model_kwargs: dict - if not None, a dict of extra keyword arguments to pass to the model. This can be used for conditioning. + + Returns: + dict: a dict with the following keys: + - 'mean': tensor - the model mean output. + - 'variance': tensor - the model variance output. + - 'log_variance': tensor - the log of 'variance'. + - 'pred_xstart': tensor - the prediction for x_0. + - 'extra': None or any type - extra information. + + Raises: + AssertionError: If the shape of t is not (B,) where B is the batch size. + AssertionError: If the shape of model_mean, model_log_variance, pred_xstart, and x is not equal. """ if model_kwargs is None: model_kwargs = {} @@ -308,6 +394,15 @@ def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, mod model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) def process_xstart(x): + """ Process the input tensor 'x' by denoising and clipping if required. + + Args: + x (tensor): Input tensor to be processed. + + Returns: + tensor: Processed tensor after denoising and clipping if required. + """ + if denoised_fn is not None: x = denoised_fn(x) if clip_denoised: @@ -332,6 +427,20 @@ def process_xstart(x): } def _predict_xstart_from_eps(self, x_t, t, eps): + """ Predict the starting point from the given state and noise. + + This function predicts the starting point from the given state and noise using the following formula: + sqrt_recip_alphas_cumprod[t] * x_t - sqrt_recipm1_alphas_cumprod[t] * eps + + Args: + x_t (tensor): The state at time t. + t (int): The time index. + eps (tensor): The noise tensor. + + Returns: + tensor: The predicted starting point. + """ + assert x_t.shape == eps.shape return ( _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t @@ -339,29 +448,60 @@ def _predict_xstart_from_eps(self, x_t, t, eps): ) def _predict_eps_from_xstart(self, x_t, t, pred_xstart): + """ Predicts the epsilon value from the starting value. + + This function calculates the epsilon value based on the input x_t, time t, and predicted starting value pred_xstart. + + Args: + x_t (tensor): The input tensor. + t (int): The time value. + pred_xstart (tensor): The predicted starting value. + + Returns: + tensor: The predicted epsilon value. + """ + return ( _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): - """ - Compute the mean for the previous step, given a function cond_fn that + """ Compute the mean for the previous step, given a function cond_fn that computes the gradient of a conditional log probability with respect to x. In particular, cond_fn computes grad(log(p(y|x))), and we want to condition on y. This uses the conditioning strategy from Sohl-Dickstein et al. (2015). + + Args: + cond_fn (function): A function that computes the gradient of a conditional log probability with respect to x. + p_mean_var (dict): A dictionary containing 'mean' and 'variance' keys representing mean and variance values. + x: The input data. + t: The time step. + model_kwargs (dict?): Additional keyword arguments for the model function. + + Returns: + torch.Tensor: The new mean computed for the previous step. """ gradient = cond_fn(x, t, **model_kwargs) new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() return new_mean def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): - """ - Compute what the p_mean_variance output would have been, should the + """ Compute what the p_mean_variance output would have been, should the model's score function be conditioned by cond_fn. See condition_mean() for details on cond_fn. Unlike condition_mean(), this instead uses the conditioning strategy from Song et al (2020). + + Args: + cond_fn (function): The conditioning function to be applied to the model's score function. + p_mean_var (dict): A dictionary containing the mean and variance of the predictive distribution. + x (tensor): The input tensor. + t (tensor): The time tensor. + model_kwargs (dict?): Additional keyword arguments for the model. Defaults to None. + + Returns: + dict: A dictionary containing the updated predictive distribution mean and variance. """ alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) @@ -383,21 +523,24 @@ def p_sample( cond_fn=None, model_kwargs=None, ): - """ - Sample x_{t-1} from the model at the given timestep. - :param model: the model to sample from. - :param x: the current tensor at x_{t-1}. - :param t: the value of t, starting at 0 for the first diffusion step. - :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. - :param denoised_fn: if not None, a function which applies to the - x_start prediction before it is used to sample. - :param cond_fn: if not None, this is a gradient function that acts - similarly to the model. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :return: a dict containing the following keys: - - 'sample': a random sample from the model. - - 'pred_xstart': a prediction of x_0. + """ Sample x_{t-1} from the model at the given timestep. + + Args: + model: object, the model to sample from. + x: tensor, the current tensor at x_{t-1}. + t: int, the value of t, starting at 0 for the first diffusion step. + clip_denoised: bool, if True, clip the x_start prediction to [-1, 1]. + denoised_fn: function, if not None, a function which applies to the + x_start prediction before it is used to sample. + cond_fn: function, if not None, this is a gradient function that acts + similarly to the model. + model_kwargs: dict, if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + + Returns: + dict: a dict containing the following keys: + - 'sample': tensor, a random sample from the model. + - 'pred_xstart': tensor, a prediction of x_0. """ out = self.p_mean_variance( model, @@ -428,23 +571,21 @@ def p_sample_loop( device=None, progress=False, ): - """ - Generate samples from the model. - :param model: the model module. - :param shape: the shape of the samples, (N, C, H, W). - :param noise: if specified, the noise from the encoder to sample. - Should be of the same shape as `shape`. - :param clip_denoised: if True, clip x_start predictions to [-1, 1]. - :param denoised_fn: if not None, a function which applies to the - x_start prediction before it is used to sample. - :param cond_fn: if not None, this is a gradient function that acts - similarly to the model. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :param device: if specified, the device to create the samples on. - If not specified, use a model parameter's device. - :param progress: if True, show a tqdm progress bar. - :return: a non-differentiable batch of samples. + """ Generate samples from the model. + + Args: + model: the model module. + shape: the shape of the samples, (N, C, H, W). + noise: if specified, the noise from the encoder to sample. Should be of the same shape as `shape`. + clip_denoised: if True, clip x_start predictions to [-1, 1]. + denoised_fn: if not None, a function which applies to the x_start prediction before it is used to sample. + cond_fn: if not None, this is a gradient function that acts similarly to the model. + model_kwargs: if not None, a dict of extra keyword arguments to pass to the model. This can be used for conditioning. + device: if specified, the device to create the samples on. If not specified, use a model parameter's device. + progress: if True, show a tqdm progress bar. + + Returns: + dict: a non-differentiable batch of samples. """ final = None for sample in self.p_sample_loop_progressive( @@ -473,12 +614,27 @@ def p_sample_loop_progressive( device=None, progress=False, ): - """ - Generate samples from the model and yield intermediate samples from + """ Generate samples from the model and yield intermediate samples from each timestep of diffusion. + Arguments are the same as p_sample_loop(). - Returns a generator over dicts, where each dict is the return value of - p_sample(). + + Args: + model: The model used for generating samples. + shape (tuple or list): The shape of the samples to be generated. + noise (Tensor?): The noise tensor for generating samples. Defaults to None. + clip_denoised (bool?): Whether to clip the denoised samples. Defaults to True. + denoised_fn (function?): The denoising function. Defaults to None. + cond_fn (function?): The conditioning function. Defaults to None. + model_kwargs (dict?): Additional keyword arguments for the model. Defaults to None. + device (str?): The device on which the operations will be performed. Defaults to None. + progress (bool?): Whether to show progress. Defaults to False. + + Returns: + generator: A generator over dicts, where each dict is the return value of p_sample(). + + Raises: + AssertionError: If the shape is not a tuple or a list. """ if device is None: device = next(model.parameters()).device @@ -521,9 +677,25 @@ def ddim_sample( model_kwargs=None, eta=0.0, ): - """ - Sample x_{t-1} from the model using DDIM. - Same usage as p_sample(). + """ Sample x_{t-1} from the model using DDIM. + + This function samples x_{t-1} from the model using DDIM (Denoising Diffusion Implicit Models). + It first calculates the mean and variance of the model's output and then conditions the score if a condition function is provided. + It then calculates the noise and mean prediction based on the model's output and other parameters. + + Args: + self: The object instance. + model: The model used for sampling. + x: The input data. + t: The time step. + clip_denoised (bool?): Whether to clip the denoised output. Defaults to True. + denoised_fn (function?): The denoising function. Defaults to None. + cond_fn (function?): The conditioning function. Defaults to None. + model_kwargs (dict?): Additional keyword arguments for the model. Defaults to None. + eta (float?): The eta value for calculating sigma. Defaults to 0.0. + + Returns: + dict: A dictionary containing the sampled data and the predicted xstart. """ out = self.p_mean_variance( model, @@ -570,8 +742,24 @@ def ddim_reverse_sample( model_kwargs=None, eta=0.0, ): - """ - Sample x_{t+1} from the model using DDIM reverse ODE. + """ Sample x_{t+1} from the model using DDIM reverse ODE. + + Args: + self: The class instance. + model: The model used for sampling. + x: The input data. + t: The time step. + clip_denoised (bool?): Whether to clip the denoised output. Defaults to True. + denoised_fn (function?): The denoising function. Defaults to None. + cond_fn (function?): The conditioning function. Defaults to None. + model_kwargs (dict?): Additional keyword arguments for the model. Defaults to None. + eta (float?): The value of eta. Defaults to 0.0. + + Returns: + dict: A dictionary containing the sampled data and the predicted start data. + + Raises: + AssertionError: If eta is not equal to 0.0. """ assert eta == 0.0, "Reverse ODE only for deterministic path" out = self.p_mean_variance( @@ -610,9 +798,24 @@ def ddim_sample_loop( progress=False, eta=0.0, ): - """ - Generate samples from the model using DDIM. + """ Generate samples from the model using DDIM. + Same usage as p_sample_loop(). + + Args: + model: The model used for generating samples. + shape: The shape of the samples to be generated. + noise (optional): The noise to be added to the samples. + clip_denoised (bool): A flag to indicate whether to clip the denoised samples. + denoised_fn: The function used for denoising. + cond_fn: The conditional function. + model_kwargs: Additional keyword arguments for the model. + device: The device to be used for generating samples. + progress (bool): A flag to indicate whether to show the progress. + eta (float): The value of eta. + + Returns: + The generated samples. """ final = None for sample in self.ddim_sample_loop_progressive( @@ -643,10 +846,26 @@ def ddim_sample_loop_progressive( progress=False, eta=0.0, ): - """ - Use DDIM to sample from the model and yield intermediate samples from + """ Use DDIM to sample from the model and yield intermediate samples from each timestep of DDIM. - Same usage as p_sample_loop_progressive(). + + Args: + model: The model used for sampling. + shape (tuple or list): The shape of the input. + noise (Tensor?): The noise tensor. Defaults to None. + clip_denoised (bool?): Whether to clip the denoised output. Defaults to True. + denoised_fn (function?): The denoising function. Defaults to None. + cond_fn (function?): The conditioning function. Defaults to None. + model_kwargs (dict?): Additional keyword arguments for the model. Defaults to None. + device (str?): The device to be used. Defaults to None. + progress (bool?): Whether to show progress. Defaults to False. + eta (float?): The eta value. Defaults to 0.0. + + Yields: + dict: Intermediate samples from each timestep of DDIM. + + Raises: + AssertionError: If the shape is not a tuple or list. """ if device is None: device = next(model.parameters()).device @@ -682,13 +901,24 @@ def ddim_sample_loop_progressive( 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. + """ Get a term for the variational lower-bound. The resulting units are bits (rather than nats, as one might expect). - This allows for comparison to other papers. - :return: a dict with the following keys: - - 'output': a shape [N] tensor of NLLs or KLs. - - 'pred_xstart': the x_0 predictions. + + Args: + model: The model used for prediction. + x_start: The initial input data. + x_t: The input data at time t. + t: The time step. + clip_denoised: A boolean indicating whether to clip the denoised output. Default is True. + model_kwargs: Additional keyword arguments for the model. Default is None. + + Returns: + dict: A dictionary with the following keys: + - 'output': A shape [N] tensor of negative log likelihoods (NLLs) or Kullback-Leibler divergences (KLs). + - 'pred_xstart': The x_0 predictions. + + Raises: + AssertionError: If the shape of decoder_nll does not match the shape of x_start. """ true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( x_start=x_start, x_t=x_t, t=t @@ -713,16 +943,22 @@ def _vb_terms_bpd( return {"output": output, "pred_xstart": out["pred_xstart"]} def training_losses(self, model, x_start, t, model_kwargs=None, noise=None): - """ - Compute training losses for a single timestep. - :param model: the model to evaluate loss on. - :param x_start: the [N x C x ...] tensor of inputs. - :param t: a batch of timestep indices. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :param noise: if specified, the specific Gaussian noise to try to remove. - :return: a dict with the key "loss" containing a tensor of shape [N]. - Some mean or variance settings may also have other keys. + """ Compute training losses for a single timestep. + + Args: + model: object, the model to evaluate loss on. + x_start: tensor, the [N x C x ...] tensor of inputs. + t: tensor, a batch of timestep indices. + model_kwargs: dict, if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + noise: tensor, if specified, the specific Gaussian noise to try to remove. + + Returns: + dict: a dict with the key "loss" containing a tensor of shape [N]. + Some mean or variance settings may also have other keys. + + Raises: + NotImplementedError: If the loss type is not supported. """ if model_kwargs is None: model_kwargs = {} @@ -787,12 +1023,15 @@ def training_losses(self, model, x_start, t, model_kwargs=None, noise=None): return terms def _prior_bpd(self, x_start): - """ - Get the prior KL term for the variational lower-bound, measured in - bits-per-dim. + """ Get the prior KL term for the variational lower-bound, measured in bits-per-dim. + This term can't be optimized, as it only depends on the encoder. - :param x_start: the [N x C x ...] tensor of inputs. - :return: a batch of [N] KL values (in bits), one per batch element. + + Args: + x_start (tensor): The [N x C x ...] tensor of inputs. + + Returns: + tensor: A batch of [N] KL values (in bits), one per batch element. """ batch_size = x_start.shape[0] t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) @@ -803,20 +1042,22 @@ def _prior_bpd(self, x_start): return mean_flat(kl_prior) / np.log(2.0) def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None): - """ - Compute the entire variational lower-bound, measured in bits-per-dim, + """ Compute the entire variational lower-bound, measured in bits-per-dim, as well as other related quantities. - :param model: the model to evaluate loss on. - :param x_start: the [N x C x ...] tensor of inputs. - :param clip_denoised: if True, clip denoised samples. - :param model_kwargs: if not None, a dict of extra keyword arguments to - pass to the model. This can be used for conditioning. - :return: a dict containing the following keys: - - total_bpd: the total variational lower-bound, per batch element. - - prior_bpd: the prior term in the lower-bound. - - vb: an [N x T] tensor of terms in the lower-bound. - - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep. - - mse: an [N x T] tensor of epsilon MSEs for each timestep. + + Args: + model: The model to evaluate loss on. + x_start: The [N x C x ...] tensor of inputs. + clip_denoised: If True, clip denoised samples. + model_kwargs: If not None, a dict of extra keyword arguments to pass to the model. This can be used for conditioning. + + Returns: + dict: A dictionary containing the following keys: + - total_bpd: The total variational lower-bound, per batch element. + - prior_bpd: The prior term in the lower-bound. + - vb: An [N x T] tensor of terms in the lower-bound. + - xstart_mse: An [N x T] tensor of x_0 MSEs for each timestep. + - mse: An [N x T] tensor of epsilon MSEs for each timestep. """ device = x_start.device batch_size = x_start.shape[0] @@ -859,13 +1100,15 @@ def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None): def _extract_into_tensor(arr, timesteps, broadcast_shape): - """ - Extract values from a 1-D numpy array for a batch of indices. - :param arr: the 1-D numpy array. - :param timesteps: a tensor of indices into the array to extract. - :param broadcast_shape: a larger shape of K dimensions with the batch - dimension equal to the length of timesteps. - :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. + """ Extract values from a 1-D numpy array for a batch of indices. + + Args: + arr (numpy.ndarray): The 1-D numpy array. + timesteps (torch.Tensor): A tensor of indices into the array to extract. + broadcast_shape (tuple): A larger shape of K dimensions with the batch dimension equal to the length of timesteps. + + Returns: + torch.Tensor: A tensor of shape [batch_size, 1, ...] where the shape has K dims. """ res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() while len(res.shape) < len(broadcast_shape): diff --git a/diffusion/timestep_sampler.py b/diffusion/timestep_sampler.py index a3f36984..77c070ea 100644 --- a/diffusion/timestep_sampler.py +++ b/diffusion/timestep_sampler.py @@ -11,10 +11,17 @@ def create_named_schedule_sampler(name, diffusion): - """ - Create a ScheduleSampler from a library of pre-defined samplers. - :param name: the name of the sampler. - :param diffusion: the diffusion object to sample for. + """ Create a ScheduleSampler from a library of pre-defined samplers. + + Args: + name (str): The name of the sampler. + diffusion (object): The diffusion object to sample for. + + Returns: + ScheduleSampler: An instance of ScheduleSampler based on the given name. + + Raises: + NotImplementedError: If the given name does not match any pre-defined samplers. """ if name == "uniform": return UniformSampler(diffusion) @@ -36,19 +43,26 @@ class ScheduleSampler(ABC): @abstractmethod def weights(self): - """ - Get a numpy array of weights, one per diffusion step. + """ Get a numpy array of weights, one per diffusion step. The weights needn't be normalized, but must be positive. + + Returns: + numpy.ndarray: An array of weights, one per diffusion step. """ def sample(self, batch_size, device): - """ - Importance-sample timesteps for a batch. - :param batch_size: the number of timesteps. - :param device: the torch device to save to. - :return: a tuple (timesteps, weights): - - timesteps: a tensor of timestep indices. - - weights: a tensor of weights to scale the resulting losses. + """ Importance-sample timesteps for a batch. + + This function importance-samples timesteps for a batch using the given batch size and torch device. + + Args: + batch_size (int): The number of timesteps. + device (torch.device): The torch device to save to. + + Returns: + tuple: A tuple containing: + - timesteps (torch.Tensor): A tensor of timestep indices. + - weights (torch.Tensor): A tensor of weights to scale the resulting losses. """ w = self.weights() p = w / np.sum(w) @@ -61,23 +75,37 @@ def sample(self, batch_size, device): class UniformSampler(ScheduleSampler): def __init__(self, diffusion): + """ Initialize the diffusion object with the given parameters. + + Args: + diffusion: The diffusion object to be initialized. + """ + self.diffusion = diffusion self._weights = np.ones([diffusion.num_timesteps]) def weights(self): + """ Returns the weights associated with the object. + + Returns: + object: The weights associated with the object. + """ + return self._weights class LossAwareSampler(ScheduleSampler): def update_with_local_losses(self, local_ts, local_losses): - """ - Update the reweighting using losses from a model. + """ Update the reweighting using losses from a model. + Call this method from each rank with a batch of timesteps and the - corresponding losses for each of those timesteps. - This method will perform synchronization to make sure all of the ranks - maintain the exact same reweighting. - :param local_ts: an integer Tensor of timesteps. - :param local_losses: a 1D Tensor of losses. + corresponding losses for each of those timesteps. This method will perform + synchronization to make sure all of the ranks maintain the exact same + reweighting. + + Args: + local_ts (Tensor): An integer Tensor of timesteps. + local_losses (Tensor): A 1D Tensor of losses. """ batch_sizes = [ th.tensor([0], dtype=th.int32, device=local_ts.device) @@ -104,21 +132,30 @@ def update_with_local_losses(self, local_ts, local_losses): @abstractmethod def update_with_all_losses(self, ts, losses): - """ - Update the reweighting using losses from a model. + """ Update the reweighting using losses from a model. + Sub-classes should override this method to update the reweighting - using losses from the model. - This method directly updates the reweighting without synchronizing - between workers. It is called by update_with_local_losses from all + using losses from the model. This method directly updates the reweighting + without synchronizing between workers. It is called by update_with_local_losses from all ranks with identical arguments. Thus, it should have deterministic behavior to maintain state across workers. - :param ts: a list of int timesteps. - :param losses: a list of float losses, one per timestep. + + Args: + ts (list): A list of int timesteps. + losses (list): A list of float losses, one per timestep. """ class LossSecondMomentResampler(LossAwareSampler): def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): + """ Initialize the object with diffusion, history_per_term, and uniform_prob. + + Args: + diffusion: The diffusion object. + history_per_term (int): Number of history per term. + uniform_prob (float): The probability of uniform distribution. + """ + self.diffusion = diffusion self.history_per_term = history_per_term self.uniform_prob = uniform_prob @@ -128,6 +165,15 @@ def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) def weights(self): + """ Calculate the weights for each timestep based on the loss history. + + If the model is not warmed up, it returns an array of ones with the same length as the number of timesteps. + Otherwise, it calculates the weights based on the root mean square of the loss history and normalizes them. + + Returns: + numpy.ndarray: An array of weights for each timestep. + """ + 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)) @@ -137,6 +183,17 @@ def weights(self): return weights def update_with_all_losses(self, ts, losses): + """ Update the loss history with all losses for given time steps. + + This method updates the loss history with all losses for the given time steps. + If the number of losses for a time step reaches the maximum history per term, + the oldest loss term is shifted out. + + Args: + ts (list): A list of time steps. + losses (list): A list of loss values corresponding to the time steps. + """ + for t, loss in zip(ts, losses): if self._loss_counts[t] == self.history_per_term: # Shift out the oldest loss term. @@ -147,4 +204,12 @@ def update_with_all_losses(self, ts, losses): self._loss_counts[t] += 1 def _warmed_up(self): + """ Check if the model has been warmed up. + + Returns True if the loss counts are equal to the history per term for all terms, indicating that the model has been warmed up. + + Returns: + bool: True if the model has been warmed up, False otherwise. + """ + return (self._loss_counts == self.history_per_term).all() diff --git a/download.py b/download.py index de22d459..cdbb5a9d 100644 --- a/download.py +++ b/download.py @@ -16,8 +16,16 @@ def find_model(model_name): - """ - Finds a pre-trained DiT model, downloading it if necessary. Alternatively, loads a model from a local path. + """ Finds a pre-trained DiT model, downloading it if necessary. Alternatively, loads a model from a local path. + + Args: + model_name (str): The name of the model or the path to the local model checkpoint. + + Returns: + dict or torch.nn.Module: The pre-trained model if found, or the loaded model checkpoint. + + Raises: + AssertionError: If the model_name is not found in the pre-trained models and is also not a valid local path. """ if model_name in pretrained_models: # Find/download our pre-trained DiT checkpoints return download_model(model_name) @@ -30,8 +38,17 @@ def find_model(model_name): def download_model(model_name): - """ - Downloads a pre-trained DiT model from the web. + """ Downloads a pre-trained DiT model from the web. + + Args: + model_name (str): The name of the pre-trained model to be downloaded. + + Returns: + torch.nn.Module: The downloaded pre-trained model. + + Raises: + AssertionError: If the specified model_name is not in the list of pretrained_models. + FileNotFoundError: If the local_path does not exist. """ assert model_name in pretrained_models local_path = f'pretrained_models/{model_name}' diff --git a/models.py b/models.py index c90eeba7..4ec66fae 100644 --- a/models.py +++ b/models.py @@ -17,6 +17,19 @@ def modulate(x, shift, scale): + """ Modulate the input tensor by applying shift and scale. + + Modulates the input tensor 'x' by adding the shift tensor and multiplying by the scale tensor. + + Args: + x (tensor): Input tensor to be modulated. + shift (tensor): Tensor representing the shift to be added. + scale (tensor): Tensor representing the scale to be applied. + + Returns: + tensor: The modulated tensor. + """ + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) @@ -29,6 +42,13 @@ class TimestepEmbedder(nn.Module): Embeds scalar timesteps into vector representations. """ def __init__(self, hidden_size, frequency_embedding_size=256): + """ Initialize the object with given hidden size and frequency embedding size. + + Args: + hidden_size (int): The size of the hidden layer. + frequency_embedding_size (int?): The size of the frequency embedding. Defaults to 256. + """ + super().__init__() self.mlp = nn.Sequential( nn.Linear(frequency_embedding_size, hidden_size, bias=True), @@ -39,13 +59,17 @@ def __init__(self, hidden_size, frequency_embedding_size=256): @staticmethod def timestep_embedding(t, dim, max_period=10000): - """ - Create sinusoidal timestep embeddings. - :param t: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param dim: the dimension of the output. - :param max_period: controls the minimum frequency of the embeddings. - :return: an (N, D) Tensor of positional embeddings. + """ Create sinusoidal timestep embeddings. + + This function creates sinusoidal timestep embeddings based on the input parameters. + + Args: + t (Tensor): A 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): The dimension of the output. + max_period (int?): Controls the minimum frequency of the embeddings. Defaults to 10000. + + Returns: + Tensor: An (N, D) Tensor of positional embeddings. """ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py half = dim // 2 @@ -59,6 +83,15 @@ def timestep_embedding(t, dim, max_period=10000): return embedding def forward(self, t): + """ Forward pass of the model to calculate the time embedding. + + Args: + t (tensor): Input tensor representing time. + + Returns: + tensor: The time embedding tensor. + """ + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) t_emb = self.mlp(t_freq) return t_emb @@ -69,6 +102,14 @@ class LabelEmbedder(nn.Module): Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance. """ def __init__(self, num_classes, hidden_size, dropout_prob): + """ Initialize the model with the specified parameters. + + Args: + num_classes (int): The number of classes for classification. + hidden_size (int): The size of the hidden layer. + dropout_prob (float): The probability of dropout. + """ + super().__init__() use_cfg_embedding = dropout_prob > 0 self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size) @@ -76,8 +117,18 @@ def __init__(self, num_classes, hidden_size, dropout_prob): self.dropout_prob = dropout_prob def token_drop(self, labels, force_drop_ids=None): - """ - Drops labels to enable classifier-free guidance. + """ Drops labels to enable classifier-free guidance. + + This function drops labels to enable classifier-free guidance. If force_drop_ids is not provided, it generates drop_ids based on a random probability. + If force_drop_ids is provided, it generates drop_ids based on the values of force_drop_ids. + It then replaces the dropped labels with the number of classes and returns the modified labels. + + Args: + labels (torch.Tensor): The input labels. + force_drop_ids (torch.Tensor?): A tensor indicating whether to force drop specific labels. + + Returns: + torch.Tensor: The modified labels after dropping. """ if force_drop_ids is None: drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob @@ -87,6 +138,19 @@ def token_drop(self, labels, force_drop_ids=None): return labels def forward(self, labels, train, force_drop_ids=None): + """ Forward pass through the neural network model. + + This method takes input labels and performs a forward pass through the neural network model. If training and dropout is enabled, or if force_drop_ids is provided, it applies token dropout to the input labels before obtaining their embeddings from the embedding table. + + Args: + labels (tensor): Input labels for the forward pass. + train (bool): Flag indicating whether the model is in training mode. + force_drop_ids (list?): List of indices to force dropout on. + + Returns: + tensor: Embeddings obtained from the embedding table. + """ + use_dropout = self.dropout_prob > 0 if (train and use_dropout) or (force_drop_ids is not None): labels = self.token_drop(labels, force_drop_ids) @@ -103,6 +167,15 @@ class DiTBlock(nn.Module): A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning. """ def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs): + """ Initializes the model with the given parameters. + + Args: + hidden_size (int): The dimension of the hidden state. + num_heads (int): The number of attention heads. + mlp_ratio (float?): The ratio of the hidden size in the feedforward network. Defaults to 4.0. + **block_kwargs: Additional keyword arguments for the attention block. + """ + 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) @@ -116,6 +189,19 @@ def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs): ) def forward(self, x, c): + """ Applies the forward pass of the module. + + Modulates the input tensor `x` using the modulation parameters obtained from the conditioning tensor `c`. + The modulation parameters are used to modulate the normalized input tensor and then applied to the input tensor using attention and MLP operations. + + Args: + x (torch.Tensor): Input tensor. + c (torch.Tensor): Conditioning tensor. + + Returns: + torch.Tensor: The modulated output tensor. + """ + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa)) x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp)) @@ -127,6 +213,14 @@ class FinalLayer(nn.Module): The final layer of DiT. """ def __init__(self, hidden_size, patch_size, out_channels): + """ Initialize the parameters for the model. + + Args: + hidden_size (int): The size of the hidden layer. + patch_size (int): The size of the patch. + out_channels (int): The number of output channels. + """ + super().__init__() self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) @@ -136,6 +230,16 @@ def __init__(self, hidden_size, patch_size, out_channels): ) def forward(self, x, c): + """ Applies forward pass through the neural network with adaptive layer normalization modulation. + + Args: + x (tensor): Input tensor. + c (tensor): Modulation tensor. + + Returns: + tensor: Output tensor after forward pass. + """ + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) x = modulate(self.norm_final(x), shift, scale) x = self.linear(x) @@ -159,6 +263,21 @@ def __init__( num_classes=1000, learn_sigma=True, ): + """ Initializes a Deep into Time (DiT) model with the specified parameters. + + Args: + input_size (int): The size of the input. + patch_size (int): The size of the patch. + in_channels (int): The number of input channels. + hidden_size (int): The size of the hidden layer. + depth (int): The depth of the model. + num_heads (int): The number of attention heads. + mlp_ratio (float): The ratio of the hidden layer size to the input size. + class_dropout_prob (float): The dropout probability for the class. + num_classes (int): The number of output classes. + learn_sigma (bool): Whether to learn sigma. + """ + super().__init__() self.learn_sigma = learn_sigma self.in_channels = in_channels @@ -180,8 +299,23 @@ def __init__( self.initialize_weights() def initialize_weights(self): + """ Initialize the weights of the model. + + This function initializes the weights of the model including transformer layers, position embedding, patch embedding, + label embedding table, timestep embedding MLP, and modulation layers in DiT blocks. + + Args: + self: The model instance. + """ + # Initialize transformer layers: def _basic_init(module): + """ Initialize the weights and biases of a linear layer using Xavier uniform initialization. + + Args: + module (torch.nn.Module): The module to initialize. + """ + if isinstance(module, nn.Linear): torch.nn.init.xavier_uniform_(module.weight) if module.bias is not None: @@ -216,9 +350,16 @@ def _basic_init(module): nn.init.constant_(self.final_layer.linear.bias, 0) def unpatchify(self, x): - """ - x: (N, T, patch_size**2 * C) - imgs: (N, H, W, C) + """ Unpatchify the input tensor to reconstruct the image. + + Args: + x (tensor): Input tensor of shape (N, T, patch_size**2 * C). + + Returns: + tensor: Reconstructed images of shape (N, H, W, C). + + Raises: + AssertionError: If the product of height and width does not match the second dimension of the input tensor. """ c = self.out_channels p = self.x_embedder.patch_size[0] @@ -231,11 +372,15 @@ def unpatchify(self, x): return imgs def forward(self, x, t, y): - """ - Forward pass of DiT. - x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) - t: (N,) tensor of diffusion timesteps - y: (N,) tensor of class labels + """ Forward pass of DiT. + + Args: + x (torch.Tensor): (N, C, H, W) tensor of spatial inputs (images or latent representations of images) + t (torch.Tensor): (N,) tensor of diffusion timesteps + y (torch.Tensor): (N,) tensor of class labels + + Returns: + torch.Tensor: (N, out_channels, H, W) tensor representing the output of the forward pass """ x = self.x_embedder(x) + self.pos_embed # (N, T, D), where T = H * W / patch_size ** 2 t = self.t_embedder(t) # (N, D) @@ -248,8 +393,18 @@ def forward(self, x, t, y): return x def forward_with_cfg(self, x, t, y, cfg_scale): - """ - Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance. + """ Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance. + + This method performs the forward pass of DiT while also batching the unconditional forward pass for classifier-free guidance. + + Args: + x (torch.Tensor): Input tensor. + t (torch.Tensor): Target tensor. + y (torch.Tensor): Output tensor. + cfg_scale (float): Scale factor for conditional guidance. + + Returns: + torch.Tensor: Concatenated tensor of conditional and unconditional guidance. """ # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb half = x[: len(x) // 2] @@ -272,10 +427,18 @@ def forward_with_cfg(self, x, t, y, cfg_scale): # https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0): - """ - grid_size: int of the grid height and width - return: - pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) + """ Generate 2D sine/cosine positional embeddings for a grid. + + This function takes the embedding dimension, grid size, and optional parameters to generate 2D sine/cosine positional embeddings for a grid. It first creates a grid using the provided grid size, then calculates the sine/cosine positional embeddings based on the grid. + + Args: + embed_dim (int): The dimension of the positional embeddings. + grid_size (int): The height and width of the grid. + cls_token (bool?): Whether to include a cls_token. Defaults to False. + extra_tokens (int?): The number of extra tokens to include. Defaults to 0. + + Returns: + numpy.ndarray: The 2D sine/cosine positional embeddings for the grid. """ grid_h = np.arange(grid_size, dtype=np.float32) grid_w = np.arange(grid_size, dtype=np.float32) @@ -290,6 +453,19 @@ def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens= def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): + """ Get 2D sine-cosine positional embeddings from a grid. + + This function takes the embedding dimension and a 2D grid as input and returns the positional embeddings + for the grid using sine and cosine functions. + + Args: + embed_dim (int): The embedding dimension, must be divisible by 2. + grid (tuple): A tuple containing the 2D grid dimensions. + + Returns: + numpy.ndarray: The positional embeddings for the input grid. + """ + assert embed_dim % 2 == 0 # use half of dimensions to encode grid_h @@ -301,10 +477,17 @@ def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): - """ - embed_dim: output dimension for each position - pos: a list of positions to be encoded: size (M,) - out: (M, D) + """ Encode positions into a 1D sinusoidal-cosine positional embedding grid. + + Args: + embed_dim (int): Output dimension for each position. + pos (ndarray): A list of positions to be encoded: size (M,). + + Returns: + ndarray: The encoded positional embedding grid of shape (M, D), where D is the output dimension. + + Raises: + AssertionError: If embed_dim is not divisible by 2. """ assert embed_dim % 2 == 0 omega = np.arange(embed_dim // 2, dtype=np.float64) @@ -326,39 +509,167 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): ################################################################################# def DiT_XL_2(**kwargs): + """ DiT model with extra large configuration and patch size of 2. + + This function returns a DiT model with a depth of 28, hidden size of 1152, patch size of 2, and 16 attention heads. + + Args: + **kwargs: Additional keyword arguments for the DiT model. + + Returns: + DiT: A DiT model with the specified configuration. + """ + return DiT(depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs) def DiT_XL_4(**kwargs): + """ Create a DiT (Diversity Transformer) model with extra large configuration using custom keyword arguments. + + This function creates a Diversity Transformer model with extra large configuration by setting the depth to 28, hidden_size to 1152, patch_size to 4, and num_heads to 16, along with any additional custom keyword arguments provided. + + Args: + **kwargs: Additional custom keyword arguments for configuring the model. + + Returns: + DiT: A Diversity Transformer model with extra large configuration. + """ + return DiT(depth=28, hidden_size=1152, patch_size=4, num_heads=16, **kwargs) def DiT_XL_8(**kwargs): + """ Create a DiT (Data-efficient image Transformer) model with extra large configuration. + + This function creates a DiT model with the specified parameters, including depth, hidden size, patch size, and number of heads. + + Args: + **kwargs: Additional keyword arguments for the DiT model. + + Returns: + DiT: A Data-efficient image Transformer model with extra large configuration. + """ + return DiT(depth=28, hidden_size=1152, patch_size=8, num_heads=16, **kwargs) def DiT_L_2(**kwargs): + """ Create a DiT model with a patch size of 2. + + This function creates a Distransformer (DiT) model with a specified depth, hidden size, patch size, and number of heads. + + Args: + **kwargs: Additional keyword arguments for the DiT model. + + Returns: + DiT: A Distransformer model with the specified parameters. + """ + return DiT(depth=24, hidden_size=1024, patch_size=2, num_heads=16, **kwargs) def DiT_L_4(**kwargs): + """ Create a DiT model with a depth of 24, hidden size of 1024, patch size of 4, and 16 attention heads. + + Args: + **kwargs: Additional keyword arguments for the DiT model. + + Returns: + DiT: An instance of the DiT model. + """ + return DiT(depth=24, hidden_size=1024, patch_size=4, num_heads=16, **kwargs) def DiT_L_8(**kwargs): + """ Create a DiT model with a patch size of 8. + + This function creates a Distransformer (DiT) model with a specified depth, hidden size, patch size, and number of heads. + + Args: + **kwargs: Additional keyword arguments for the DiT model. + + Returns: + torch.nn.Module: A DiT model with the specified parameters. + """ + return DiT(depth=24, hidden_size=1024, patch_size=8, num_heads=16, **kwargs) def DiT_B_2(**kwargs): + """ Create a DiT model with specific configuration. + + This function creates a Distransformer (DiT) model with the specified configuration parameters. + + Args: + **kwargs: Additional keyword arguments for configuring the DiT model. + + Returns: + DiT: A Distransformer model with the specified configuration. + """ + return DiT(depth=12, hidden_size=768, patch_size=2, num_heads=12, **kwargs) def DiT_B_4(**kwargs): + """ Create a Diverse Transformer model with a patch size of 4. + + This function creates a Diverse Transformer model with the specified depth, hidden size, patch size, and number of heads. + + Args: + **kwargs: Additional keyword arguments for the DiT model. + + Returns: + DiT: A Diverse Transformer model with the specified configuration. + """ + return DiT(depth=12, hidden_size=768, patch_size=4, num_heads=12, **kwargs) def DiT_B_8(**kwargs): + """ Create a DiT model with a depth of 12, hidden size of 768, patch size of 8, and 12 attention heads. + + Args: + **kwargs: Additional keyword arguments for the DiT model. + + Returns: + DiT: A DiT model with the specified configuration. + """ + return DiT(depth=12, hidden_size=768, patch_size=8, num_heads=12, **kwargs) def DiT_S_2(**kwargs): + """ Create a DiT model with specific configuration. + + This function creates a DiT (Data-efficient image Transformer) model with the specified configuration parameters. + + Args: + **kwargs: Additional keyword arguments for configuring the DiT model. + + Returns: + DiT: A Data-efficient image Transformer model. + """ + return DiT(depth=12, hidden_size=384, patch_size=2, num_heads=6, **kwargs) def DiT_S_4(**kwargs): + """ Create a DiT model with specific configuration settings. + + This function creates a Distransformer (DiT) model with the specified depth, hidden size, patch size, and number of heads. + + Args: + **kwargs: Additional keyword arguments for configuring the DiT model. + + Returns: + DiT: A Distransformer model with the specified configuration settings. + """ + return DiT(depth=12, hidden_size=384, patch_size=4, num_heads=6, **kwargs) def DiT_S_8(**kwargs): + """ Create a DiT model with specific configuration parameters. + + This function creates a Diverse Transformer (DiT) model with the specified configuration parameters, including depth, hidden size, patch size, and number of heads. + + Args: + **kwargs: Additional keyword arguments for configuring the DiT model. + + Returns: + DiT: A Diverse Transformer model with the specified configuration. + """ + return DiT(depth=12, hidden_size=384, patch_size=8, num_heads=6, **kwargs) diff --git a/sample.py b/sample.py index a4152afd..10cc2562 100644 --- a/sample.py +++ b/sample.py @@ -19,6 +19,18 @@ def main(args): + """ Setup PyTorch, load a model, create sampling noise, and sample images. + + This function sets up PyTorch, loads a model, creates sampling noise, and samples images using the provided arguments. + + Args: + args (argparse.Namespace): The input arguments for the function. + + + Raises: + AssertionError: If the specified model, image size, or number of classes do not match the expected values. + """ + # Setup PyTorch: torch.manual_seed(args.seed) torch.set_grad_enabled(False) diff --git a/sample_ddp.py b/sample_ddp.py index 0a6b1ab6..ca8650ba 100644 --- a/sample_ddp.py +++ b/sample_ddp.py @@ -26,8 +26,17 @@ def create_npz_from_sample_folder(sample_dir, num=50_000): - """ - Builds a single .npz file from a folder of .png samples. + """ Builds a single .npz file from a folder of .png samples. + + Args: + sample_dir (str): The directory path containing the .png samples. + num (int): The number of samples to be included in the .npz file. Default is 50,000. + + Returns: + str: The file path of the created .npz file. + + Raises: + AssertionError: If the shape of the samples array does not match the expected shape. """ samples = [] for i in tqdm(range(num), desc="Building .npz file from samples"): @@ -43,8 +52,16 @@ def create_npz_from_sample_folder(sample_dir, num=50_000): def main(args): - """ - Run sampling. + """ Run sampling. + + This function runs the sampling process using Distributed Data Parallel (DDP) if available. It sets up the DDP, loads the model, and generates samples on each GPU. + + Args: + args (argparse.Namespace): A namespace containing the command-line arguments. + + + Raises: + AssertionError: If sampling with DDP is attempted without any available GPU, or if certain conditions for args are not met. """ torch.backends.cuda.matmul.allow_tf32 = args.tf32 # True: fast but may lead to some small numerical differences assert torch.cuda.is_available(), "Sampling with DDP requires at least one GPU. sample.py supports CPU-only usage" diff --git a/train.py b/train.py index 7cfee808..0a6837c9 100644 --- a/train.py +++ b/train.py @@ -38,8 +38,19 @@ @torch.no_grad() def update_ema(ema_model, model, decay=0.9999): - """ - Step the EMA model towards the current model. + """ Step the Exponential Moving Average (EMA) model towards the current model. + + This function updates the EMA model by applying a decay factor to the EMA parameters and adding the current model parameters + with a specific weight. + + Args: + ema_model (torch.nn.Module): The EMA model to be updated. + model (torch.nn.Module): The current model. + decay (float?): The decay factor for updating the EMA model. Defaults to 0.9999. + + + Note: + The EMA model is updated in-place. """ ema_params = OrderedDict(ema_model.named_parameters()) model_params = OrderedDict(model.named_parameters()) @@ -50,23 +61,32 @@ def update_ema(ema_model, model, decay=0.9999): def requires_grad(model, flag=True): - """ - Set requires_grad flag for all parameters in a model. + """ Set requires_grad flag for all parameters in a model. + + Args: + model (torch.nn.Module): The model for which the requires_grad flag needs to be set. + flag (bool?): The flag to be set for requires_grad. Defaults to True. """ for p in model.parameters(): p.requires_grad = flag def cleanup(): - """ - End DDP training. + """ End DDP training. + + This function ends the DDP (Distributed Data Parallel) training by destroying the process group. """ dist.destroy_process_group() def create_logger(logging_dir): - """ - Create a logger that writes to a log file and stdout. + """ Create a logger that writes to a log file and stdout. + + Args: + logging_dir (str): The directory path where the log file will be stored. + + Returns: + logging.Logger: The logger object for writing logs. """ if dist.get_rank() == 0: # real logger logging.basicConfig( @@ -83,9 +103,17 @@ def create_logger(logging_dir): def center_crop_arr(pil_image, image_size): - """ - Center cropping implementation from ADM. - https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126 + """ Center crop an image array to the specified size. + + This function resizes the input image to ensure that the smaller dimension is at least twice the specified size. + It then scales the image to match the specified size and crops the center portion of the image. + + Args: + pil_image (PIL.Image): The input image to be cropped. + image_size (int): The desired size of the cropped image. + + Returns: + PIL.Image: The center-cropped image. """ while min(*pil_image.size) >= 2 * image_size: pil_image = pil_image.resize( @@ -108,8 +136,16 @@ def center_crop_arr(pil_image, image_size): ################################################################################# def main(args): - """ - Trains a new DiT model. + """ Trains a new DiT model. + + This function trains a new Differentiable Transformer (DiT) model using the provided arguments. It sets up the distributed training environment, initializes the experiment folder, creates the DiT model, sets up the optimizer, prepares the data, and performs the training loop. + + Args: + args (argparse.Namespace): The input arguments for training the DiT model. + + + Raises: + AssertionError: If no GPU is available for training. """ assert torch.cuda.is_available(), "Training currently requires at least one GPU."