modelzoo.vision.pytorch.dit.input.DiffusionBaseProcessor.DiffusionBaseProcessor#

class modelzoo.vision.pytorch.dit.input.DiffusionBaseProcessor.DiffusionBaseProcessor[source]#

Bases: object

Methods

check_split_valid

create_dataloader

Dataloader returns a dict with keys:

create_dataset

custom_collate_fn

process_transform

__init__(params)[source]#
create_dataloader(dataset, is_training=False)[source]#
Dataloader returns a dict with keys:

“input”: Tensor of shape (batch_size, latent_channels, latent_height, latent_width) “label”: Tensor of shape (batch_size, ) with dropout applied with label_dropout_rate “diffusion_noise”: Tensor of shape (batch_size, latent_channels, latent_height, latent_width)

represents diffusion noise to be applied

“timestep”: Tensor of shape (batch_size, ) that

indicates the timesteps for each diffusion sample