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)[source]

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)[source]

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)[source]

Generating random synthetic data.

Parameters:

  • label_info (omegaconf.dictconfig.DictConfig) – 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

variation_api(syn_data)[source]

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