cerebras.modelzoo.data.vision.diffusion.dit_transforms.NoiseGenerator#
- class cerebras.modelzoo.data.vision.diffusion.dit_transforms.NoiseGenerator(*args, **kwargs)[source]#
Bases:
torch.nn.ModuleMethods
: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