pe.api.image package

Subpackages

• pe.api.image.improved_diffusion_lib package
  • Submodules
  • pe.api.image.improved_diffusion_lib.gaussian_diffusion module
    • SkippedSpacedDiffusion
    • create_gaussian_diffusion()
  • pe.api.image.improved_diffusion_lib.unet module
    • FP32UNetModel
    • create_model()

Submodules

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 (dict) – 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.Data

variation_api(syn_data)

Generating variations of the synthetic data.

Parameters:

syn_data (pe.data.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.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.

training: bool

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)

pe.api.image.stable_diffusion_api module

class pe.api.image.stable_diffusion_api.StableDiffusion(prompt, variation_degrees, width=512, height=512, random_api_checkpoint='CompVis/stable-diffusion-v1-4', random_api_guidance_scale=7.5, random_api_num_inference_steps=50, random_api_batch_size=10, variation_api_checkpoint='CompVis/stable-diffusion-v1-4', variation_api_guidance_scale=7.5, variation_api_num_inference_steps=50, variation_api_batch_size=10)

Bases: API

The API that uses the Stable Diffusion model to generate synthetic data.

__init__(prompt, variation_degrees, width=512, height=512, random_api_checkpoint='CompVis/stable-diffusion-v1-4', random_api_guidance_scale=7.5, random_api_num_inference_steps=50, random_api_batch_size=10, variation_api_checkpoint='CompVis/stable-diffusion-v1-4', variation_api_guidance_scale=7.5, variation_api_num_inference_steps=50, variation_api_batch_size=10)

Constructor.

Parameters:

• prompt (str or dict) – The prompt used for each label name. It can be either a string or a dictionary. If it is a string, it should be the path to a JSON file that contains the prompt for each label name. If it is a dictionary, it should be a dictionary that maps each label name to its prompt

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

• width (int, optional) – The width of the generated images, defaults to 512

• height (int, optional) – The height of the generated images, defaults to 512

• random_api_checkpoint (str, optional) – The checkpoint of the random API, defaults to “CompVis/stable-diffusion-v1-4”

• random_api_guidance_scale (float, optional) – The guidance scale of the random API, defaults to 7.5

• random_api_num_inference_steps (int, optional) – The number of inference steps of the random API, defaults to 50

• random_api_batch_size (int, optional) – The batch size of the random API, defaults to 10

• variation_api_checkpoint (str, optional) – The checkpoint of the variation API, defaults to “CompVis/stable-diffusion-v1-4”

• variation_api_guidance_scale (float or list[float], optional) – The guidance scale of the variation API utilized at each PE iteration. If a single float is provided, the same guidance scale will be used for all iterations. Defaults to 7.5

• variation_api_num_inference_steps (int or list[int], optional) – The number of inference steps of the variation API utilized at each PE iteration. If a single int is provided, the same number of inference steps will be used for all iterations. Defaults to 50

• variation_api_batch_size (int, optional) – The batch size of the variation API, defaults to 10

Raises:

ValueError – If the prompt is neither a string nor a dictionary

random_api(label_info, num_samples)

Generating random synthetic data.

Parameters:

• label_info (dict) – The info of the label

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

Returns:

The data object of the generated synthetic data

Return type:

pe.data.data.Data

variation_api(syn_data)

Generating variations of the synthetic data.

Parameters:

syn_data (pe.data.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.Data