pe.api.image.improved_diffusion_api module

class pe.api.image.improved_diffusion_api.ImprovedDiffusion(variation_degrees, model_path, model_image_size=64, num_channels=192, num_res_blocks=3, learn_sigma=True, class_cond=True, use_checkpoint=False, attention_resolutions='16,8', num_heads=4, num_heads_upsample=-1, use_scale_shift_norm=True, dropout=0.0, diffusion_steps=4000, sigma_small=False, noise_schedule='cosine', use_kl=False, predict_xstart=False, rescale_timesteps=False, rescale_learned_sigmas=False, timestep_respacing='100', batch_size=2000, use_ddim=True, clip_denoised=True, use_data_parallel=True)

Bases: API

The image API that utilizes improved diffusion models from https://arxiv.org/abs/2102.09672.

__init__(variation_degrees, model_path, model_image_size=64, num_channels=192, num_res_blocks=3, learn_sigma=True, class_cond=True, use_checkpoint=False, attention_resolutions='16,8', num_heads=4, num_heads_upsample=-1, use_scale_shift_norm=True, dropout=0.0, diffusion_steps=4000, sigma_small=False, noise_schedule='cosine', use_kl=False, predict_xstart=False, rescale_timesteps=False, rescale_learned_sigmas=False, timestep_respacing='100', batch_size=2000, use_ddim=True, clip_denoised=True, use_data_parallel=True)

Constructor. See https://github.com/openai/improved-diffusion for the explanation of the parameters not listed here.

Parameters:

• variation_degrees (int or list[int]) – The variation degrees utilized at each PE iteration. If a single int is provided, the same variation degree will be used for all iterations.

• model_path (str) – The path of the model checkpoint

• diffusion_steps (int, optional) – The total number of diffusion steps, defaults to 4000

• timestep_respacing (str or list[str], optional) – The step configurations for image generation utilized at each PE iteration. If a single str is provided, the same step configuration will be used for all iterations. Defaults to “100”

• batch_size (int, optional) – The batch size for image generation, defaults to 2000

• use_data_parallel (bool, optional) – Whether to use data parallel during image generation, defaults to True

random_api(label_info, num_samples)

Generating random synthetic data.

Parameters:

• label_info (omegaconf.dictconfig.DictConfig) – The info of the label, not utilized in this API

• num_samples (int) – The number of random samples to generate

Returns:

The data object of the generated synthetic data

Return type:

pe.data.Data

variation_api(syn_data)

Generating variations of the synthetic data.

Parameters:

syn_data (pe.data.Data) – The data object of the synthetic data

Returns:

The data object of the variation of the input synthetic data

Return type:

pe.data.Data

class pe.api.image.improved_diffusion_api.ImprovedDiffusion270M(variation_degrees, model_path=None, batch_size=2000, timestep_respacing='100', use_data_parallel=True)

Bases: ImprovedDiffusion

CHECKPOINT_URL = 'https://openaipublic.blob.core.windows.net/diffusion/march-2021/imagenet64_cond_270M_250K.pt'

The URL of the checkpoint path

__init__(variation_degrees, model_path=None, batch_size=2000, timestep_respacing='100', use_data_parallel=True)

The “Class-conditional ImageNet-64 model (270M parameters, trained for 250K iterations)” model from the Improved Diffusion paper.

Parameters:

• variation_degrees (list[int]) – The variation degrees utilized at each PE iteration

• model_path (str) – The path of the model checkpoint. If not provided, the checkpoint will be downloaded from the CHECKPOINT_URL

• batch_size (int, optional) – The batch size for image generation, defaults to 2000

• timestep_respacing (str, optional) – The step configuration for image generation, defaults to “100”

• use_data_parallel (bool, optional) – Whether to use data parallel during image generation, defaults to True

class pe.api.image.improved_diffusion_api.Sampler(model, diffusion)

Bases: Module

A wrapper around the model and diffusion modules that handles the entire sampling process, so as to reduce the communiation rounds between GPUs when using DataParallel.

forward(clip_denoised, model_kwargs, start_t, start_image, use_ddim, noise, image_size)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

pe.api.image.improved_diffusion_api.sample(sampler, num_samples, start_t, batch_size, use_ddim, image_size, clip_denoised, class_cond, device, start_image=None, labels=None)