peal.generators.ddpm_generator.DDPM

class peal.generators.ddpm_generator.DDPM[source]

Bases: EditCapableGenerator, InvertibleGenerator

Pixel-space denoising diffusion model used as a PEAL generator.

Wraps the guided-diffusion UNet and gaussian diffusion created by create_model_and_diffusion from the vendored ACE code. Besides plain sampling it provides the invertible interface (encode/decode), a repaint step that keeps everything outside a smooth change mask, and edit, which runs the ACE or FastDiME counterfactual attack.

Parameters:

  • config (str or DDPMConfig) – Config path or object; loaded with load_yaml_config. The image size is overwritten with the last entry of config.data.input_size.

  • model_dir (str, optional) – Run directory holding the weights, logs and outputs. Defaults to config.base_path.

  • device (str, optional) – Device the UNet is moved to.

  • predictor_dataset (optional) – Accepted for interface compatibility; unused.

Variables:

  • diffusion (model,) – The UNet and the gaussian diffusion process.

  • dataset – First dataset built from config.data; used for the normalisation between the model range and the pytorch default range.

  • model_path (str) – final.cpl in the run directory if it exists, else final.pt (downloaded from config.download_weights when missing). If no weights are found, an existing run directory is moved aside with a timestamp suffix and a fresh one is created.

  • noise_fn (callable) – torch.randn_like for stochastic sampling, torch.zeros_like otherwise.

Methods

decode

Denoise a latent back to an image with the reverse diffusion loop.

edit

Create counterfactuals for x_in with ACE or FastDiME.

encode

Map images to the noisy latent at a fraction t of the schedule.

repaint

Re-diffuse an edited image while pinning the unchanged region.

sample_x

Draw images from the prior with the full ancestral sampling loop.

train_model

Train the UNet with the guided-diffusion TrainLoop.
__init__(config, model_dir=None, device='cpu', predictor_dataset=None)[source]

Load the config, build the UNet and diffusion and restore weights.

sample_x(batch_size=None, renormalize=True)[source]

Draw images from the prior with the full ancestral sampling loop.

Parameters:

  • batch_size (int, optional) – Number of images; defaults to config.batch_size.

  • renormalize (bool, optional) – Map the samples from the model range into the pytorch default range via dataset.project_to_pytorch_default.

Returns:

Samples of shape (batch_size, *config.data.input_size).

Return type:

torch.Tensor

encode(x, t=1.0, stochastic=None, num_steps=None)[source]

Map images to the noisy latent at a fraction t of the schedule.

With stochastic == "fully" the latent is a single q_sample draw at the respaced timestep, clamped to [-1, 1]. Otherwise the image is inverted deterministically by iterating ddim_reverse_sample over the first t * timestep_respacing steps.

Parameters:

  • x (torch.Tensor) – Batch of images in the model range, shape (B, C, H, W).

  • t (float, optional) – Fraction of config.timestep_respacing to noise up to.

  • stochastic (bool or str, optional) – "fully" for one-shot noising, anything else for DDIM inversion. Defaults to config.stochastic.

  • num_steps (optional) – Accepted for interface compatibility; unused.

Returns:

Latent with the same shape as x.

Return type:

torch.Tensor

decode(z, t=1.0, stochastic=None, num_steps=None)[source]

Denoise a latent back to an image with the reverse diffusion loop.

Iterates p_mean_variance from step t * timestep_respacing - 1 down to 0, taking the posterior mean and adding the posterior noise in every step but the last when config.stochastic is set.

Parameters:

  • z (torch.Tensor or list) – Latent of shape (B, C, H, W); a one-element list is unwrapped.

  • t (float, optional) – Fraction of the respaced schedule the latent was noised to; must match the value used in encode().

  • stochastic (bool, optional) – Defaults to config.stochastic. Note that the per-step noise is gated by config.stochastic rather than by this argument.

  • num_steps (optional) – Accepted for interface compatibility; unused.

Returns:

Reconstructed images in the model range.

Return type:

torch.Tensor

repaint(x, pe, inpaint, dilation, t, stochastic, boolmask_in=None, max_avg_combination=0.5, exceptions=None)[source]

Re-diffuse an edited image while pinning the unchanged region.

A smooth change mask between x and the edited image pe is built with generate_smooth_mask; everything whose dilated mask value stays below inpaint is considered background and is, in every reverse step, replaced by the correspondingly noised original. The final image also takes the background pixels straight from x.

Parameters:

  • x (torch.Tensor) – Original images, shape (B, C, H, W).

  • pe (torch.Tensor) – Edited (“pre-explanation”) images of the same shape.

  • inpaint (float) – Mask threshold; 0 disables the per-step re-injection of the original but not the final composition.

  • dilation (float) – Dilation of the change mask, passed to generate_smooth_mask.

  • t (float) – Fraction of config.timestep_respacing to restart from.

  • stochastic (bool) – Add posterior noise during the reverse loop and noise when re-injecting the original.

  • boolmask_in (torch.Tensor, optional) – Background mask of a previous attempt; its complement is unioned into the current background mask, so already-edited regions stay editable.

  • max_avg_combination (float, optional) – Mixing weight between the max and the mean channel difference in generate_smooth_mask.

  • exceptions (torch.Tensor, optional) – Per-sample flags; where the entry is 1 the background mask is zeroed, i.e. that sample is not repainted at all.

Returns:

train_model()[source]

Train the UNet with the guided-diffusion TrainLoop.

Configures the guided-diffusion logger and the schedule sampler named by config.schedule_sampler, builds a training dataloader whose epoch length is config.max_steps, logs a batch of training images to <model_dir>/logs and runs the loop, which writes checkpoints into model_dir every config.save_interval steps. If the model was not trained before, an existing run directory is moved aside with a timestamp suffix.

Notes

config.x_selection overrides the dataset’s task_config so that the generator can be trained on a different input column than the one the dataset was built for.

edit(x_in, target_confidence_goal, source_classes, target_classes, predictor, explainer_config, predictor_datasets, pbar=None, base_path='', mode='', boolmask_in=None, attempt_number=None)[source]

Create counterfactuals for x_in with ACE or FastDiME.

Trains the diffusion model first if needed, optionally distils the predictor into a surrogate that supplies the guidance gradients (cached at <base_path>/explainer/distilled_predictor/model.cpl), and then calls the vendored ace_main or fastdime_main depending on explainer_config.subtype.

The attack is repeated explainer_config.num_attempts times. Any of attack_iterations, sampling_time_fraction and sampling_inpaint given as a two-element list is interpolated linearly between the attempts, dist_l1/dist_l2 geometrically, and the seed is incremented per attempt. All attempts are concatenated into the returned batch rather than being filtered by success.

Parameters:

  • x_in (torch.Tensor) – Images to explain, shape (B, C, H, W).

  • target_confidence_goal (float) – Accepted for interface compatibility; the stopping criterion is taken from explainer_config.

  • source_classes (torch.Tensor) – Current and desired class index per sample, shape (B,).

  • target_classes (torch.Tensor) – Current and desired class index per sample, shape (B,).

  • predictor (nn.Module) – Classifier whose decision the counterfactual has to flip; also used to score the results.

  • explainer_config (ACEConfig) – Attack hyper-parameters; subtype selects "ACE" or "FastDiME", l_perc/l_perc_layer enable the VGG perceptual loss, distilled_predictor the surrogate.

  • predictor_datasets (list) – [train, val, test]-style datasets; entry 1 is passed to the attack as the predictor’s dataset and entry set is used for distillation.

  • pbar (optional) – Progress bar, unused here.

  • base_path (str, optional) – Run directory used to cache the distilled predictor.

  • mode (str, optional) – Unused; kept for interface compatibility.

  • boolmask_in (optional) – Unused; kept for interface compatibility with other generators.

  • attempt_number (optional) – Unused; attempts are looped over internally.

Returns:

(counterfactuals, differences, target confidences, originals, histories, zeros), each with B * num_attempts entries. The differences are x_in - counterfactual and the trailing zeros stand in for the per-sample masks other generators return.

Return type:

tuple of list

Raises:

Exception – If explainer_config.subtype is neither ACE nor FastDiME.