cerebras.modelzoo.data.vision.diffusion.dit_transforms.NoiseGenerator#

class cerebras.modelzoo.data.vision.diffusion.dit_transforms.NoiseGenerator(*args, **kwargs)[source]#

Bases: torch.nn.Module

Methods

forward

:param : param input : Float tensor of size (B, C, H, W).

forward(input, label)[source]#

:param : param input : Float tensor of size (B, C, H, W). :param : param label : Int tensor of size (B, ).

Returns

A dict corresponding to the noisy images, ground truth noises and the timesteps corresponding to the scheduled noise variance with the following keys and shapes.

”input”: Tensor of shape (batch_size, C, H, W). This tensor is simply passed through. “label”: Tensor of shape (batch_size, ) representing labels. This tensor is simply passed through. “diffusion_noise”: Tensor of shape (batch_size, channels, height, width)

represents diffusion noise to be applied

”timestep”: Tensor of shape (batch_size, ) that indicates the timesteps for each diffusion sample “vae_noise”: Tensor of shape (batch_size, latent_channels, latent_height, latent_width)

represents the noise sample to be used with reparametrization of VAE