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

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

Bases: torch.nn.Module

Methods

forward

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

__init__(width, height, channels, num_diffusion_steps, seed=None)

forward(input, label)

: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

__call__(*args: Any, **kwargs: Any) → Any

Call self as a function.

static __new__(cls, *args: Any, **kwargs: Any) → Any

previous

cerebras.modelzoo.data.vision.diffusion.dit_transforms.LabelDropout

next

cerebras.modelzoo.data.vision.masked_auto_encoding