sketchkit.image2sketch.swiftsketch.diffusion package¶

Submodules¶

sketchkit.image2sketch.swiftsketch.diffusion.Loss_computation module¶

class sketchkit.image2sketch.swiftsketch.diffusion.Loss_computation.L1_points[source]¶

Bases: Module

forward(control_points, svg_control_points)[source]¶

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.

class sketchkit.image2sketch.swiftsketch.diffusion.Loss_computation.LPIPS(args, pretrained=True, normalize=True, pre_relu=True)[source]¶

Bases: Module

class _FeatureExtractor(pretrained, pre_relu)[source]¶

Bases: Module

forward(x)[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

_l2_normalize_features(x, eps=1e-10)[source]¶

forward(pred, target, mode='train')[source]¶: Compare VGG features of two inputs.

class sketchkit.image2sketch.swiftsketch.diffusion.Loss_computation.Loss(args)[source]¶

Bases: Module

forward(sketches, targets, output_control_points, target_control_points, mode='train')[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

get_losses_to_apply(args)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util module¶

Helpers to train with 16-bit precision.

class sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.MixedPrecisionTrainer(*, model, use_fp16=False, fp16_scale_growth=0.001, initial_lg_loss_scale=20.0)[source]¶

Bases: object

_compute_norms(grad_scale=1.0)[source]¶

_optimize_fp16(opt: Optimizer)[source]¶

_optimize_normal(opt: Optimizer)[source]¶

backward(loss: Tensor)[source]¶

master_params_to_state_dict(master_params)[source]¶

optimize(opt: Optimizer)[source]¶

state_dict_to_master_params(state_dict)[source]¶

zero_grad()[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.check_overflow(value)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.convert_module_to_f16(l)[source]¶: Convert primitive modules to float16.

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.convert_module_to_f32(l)[source]¶: Convert primitive modules to float32, undoing convert_module_to_f16().

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.get_param_groups_and_shapes(named_model_params)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.make_master_params(param_groups_and_shapes)[source]¶: Copy model parameters into a (differently-shaped) list of full-precision parameters.

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.master_params_to_model_params(param_groups_and_shapes, master_params)[source]¶: Copy the master parameter data back into the model parameters.

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.master_params_to_state_dict(model, param_groups_and_shapes, master_params, use_fp16)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.model_grads_to_master_grads(param_groups_and_shapes, master_params)[source]¶: Copy the gradients from the model parameters into the master parameters from make_master_params().

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.param_grad_or_zeros(param)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.state_dict_to_master_params(model, state_dict, use_fp16)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.unflatten_master_params(param_group, master_param)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.zero_grad(model_params)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.fp16_util.zero_master_grads(master_params)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.gaussian_diffusion module¶

This code started out as a PyTorch port of Ho et al’s diffusion models: https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py

class sketchkit.image2sketch.swiftsketch.diffusion.gaussian_diffusion.GaussianDiffusion(args, *, betas, model_mean_type, model_var_type, rescale_timesteps=False)[source]¶

Bases: object

Utilities for training and sampling diffusion models.

Ported directly from here, and then adapted over time to further experimentation. https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42

Parameters:

betas – a 1-D numpy array of betas for each diffusion timestep, starting at T and going to 1.
model_mean_type – a ModelMeanType determining what the model outputs.
model_var_type – a ModelVarType determining how variance is output.
rescale_timesteps – if True, pass floating point timesteps into the model so that they are always scaled like in the original paper (0 to 1000).

_predict_eps_from_xstart(x_t, t, pred_xstart)[source]¶

_predict_xstart_from_eps(x_t, t, eps)[source]¶

_predict_xstart_from_xprev(x_t, t, xprev)[source]¶

_scale_timesteps(t)[source]¶

p_mean_variance(model, x, t, clip_denoised=True, denoised_fn=None, image_features=None, scale=None)[source]¶

Apply the model to get p(x_{t-1} | x_t), as well as a prediction of the initial x, x_0.

Parameters:

model – the model, which takes a signal and a batch of timesteps as input.
x – the [N x C x …] tensor at time t.
t – a 1-D Tensor of timesteps.
clip_denoised – if True, clip the denoised signal into [-1, 1].
denoised_fn – if not None, a function which applies to the x_start prediction before it is used to sample. Applies before clip_denoised.
image_features – features used to condition the model’s prediction.

Returns:

a dict with the following keys: - ‘mean’: the model mean output. - ‘variance’: the model variance output. - ‘log_variance’: the log of ‘variance’. - ‘pred_xstart’: the prediction for x_0.

p_sample(model, x, t, clip_denoised=True, denoised_fn=None, image_features=None, scale=None, const_noise=False)[source]¶

Sample x_{t-1} from the model at the given timestep.

Parameters:

model – the model to sample from.
x – the current tensor at x_{t-1}.
t – the value of t, starting at 0 for the first diffusion step.
clip_denoised – if True, clip the x_start prediction to [-1, 1].
denoised_fn – if not None, a function which applies to the x_start prediction before it is used to sample.
image_features – features used to condition the model’s prediction.
scale – guidance_param
const_noise – If True, will noise all samples with the same noise throughout sampling

Returns:

a dict containing the following keys: - ‘sample’: a random sample from the model. - ‘pred_xstart’: a prediction of x_0.

p_sample_loop(model, shape, noise=None, clip_denoised=True, denoised_fn=None, image_features=None, scale=None, device=None, progress=False, skip_timesteps=0, init_image=None, dump_steps=None, const_noise=False)[source]¶

Generate samples from the model.

Parameters:

model – the model module.
shape – the shape of the samples, (N, C, H, W). #(args.batch_size, args.num_paths, model.npoints, model.nfeats)
noise – if specified, the noise from the encoder to sample. Should be of the same shape as shape.
clip_denoised – if True, clip x_start predictions to [-1, 1].
denoised_fn – if not None, a function which applies to the x_start prediction before it is used to sample.
image_features – features used to condition the model’s prediction.
scale – guidance_param
device – if specified, the device to create the samples on. If not specified, use a model parameter’s device.
progress – if True, show a tqdm progress bar.
const_noise – If True, will noise all samples with the same noise throughout sampling

Returns:

a non-differentiable batch of samples.

p_sample_loop_progressive(model, shape, noise=None, clip_denoised=True, denoised_fn=None, image_features=None, scale=None, device=None, progress=False, skip_timesteps=0, init_image=None, const_noise=False)[source]¶

Generate samples from the model and yield intermediate samples from each timestep of diffusion.

Arguments are the same as p_sample_loop(). Returns a generator over dicts, where each dict is the return value of p_sample().

plot_forward_pass_with_predicted_x0(model, data, device, scale)[source]¶

q_mean_variance(x_start, t)[source]¶

Get the distribution q(x_t | x_0).

Parameters:

x_start – the [N x C x …] tensor of noiseless inputs.
t – the number of diffusion steps (minus 1). Here, 0 means one step.

Returns:

A tuple (mean, variance, log_variance), all of x_start’s shape.

q_posterior_mean_variance(x_start, x_t, t)[source]¶: Compute the mean and variance of the diffusion posterior:

q(x_{t-1} | x_t, x_0)

q_sample(x_start, t, noise=None)[source]¶

Diffuse the dataset for a given number of diffusion steps.

In other words, sample from q(x_t | x_0).

Parameters:

x_start – the initial dataset batch.
t – the number of diffusion steps (minus 1). Here, 0 means one step.
noise – if specified, the split-out normal noise.

Returns:

A noisy version of x_start.

training_losses(model, x_start, x_start_randered_images, image_features, t, step, resume_step, noise=None)[source]¶

Compute training losses for a single timestep.

Parameters:

model – the model to evaluate loss on.
x_start – the [N x C x …] tensor of inputs. (control points)
x_start_randered_image – the [N x C x …] tensor of inputs_randered_images. (the sketches for image loss)
t – a batch of timestep indices.

:param image_features : a batch of image features (the condition) :param noise: if specified, the specific Gaussian noise to try to remove. :return: a dict with the key “loss” containing a tensor of shape [N].

Some mean or variance settings may also have other keys.

class sketchkit.image2sketch.swiftsketch.diffusion.gaussian_diffusion.ModelMeanType(*values)[source]¶

Bases: Enum

Which type of output the model predicts.

EPSILON = 3¶

PREVIOUS_X = 1¶

START_X = 2¶

class sketchkit.image2sketch.swiftsketch.diffusion.gaussian_diffusion.ModelVarType(*values)[source]¶

Bases: Enum

What is used as the model’s output variance.

The LEARNED_RANGE option has been added to allow the model to predict values between FIXED_SMALL and FIXED_LARGE, making its job easier.

FIXED_LARGE = 3¶

FIXED_SMALL = 2¶

LEARNED = 1¶

LEARNED_RANGE = 4¶

sketchkit.image2sketch.swiftsketch.diffusion.gaussian_diffusion._extract_into_tensor(arr, timesteps, broadcast_shape)[source]¶

Extract values from a 1-D numpy array for a batch of indices.

Parameters:

arr – the 1-D numpy array.
timesteps – a tensor of indices into the array to extract.
broadcast_shape – a larger shape of K dimensions with the batch dimension equal to the length of timesteps.

Returns:

a tensor of shape [batch_size, 1, …] where the shape has K dims.

sketchkit.image2sketch.swiftsketch.diffusion.gaussian_diffusion.betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999)[source]¶

Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of (1-beta) over time from t = [0,1].

Parameters:

num_diffusion_timesteps – the number of betas to produce.
alpha_bar – a lambda that takes an argument t from 0 to 1 and produces the cumulative product of (1-beta) up to that part of the diffusion process.
max_beta – the maximum beta to use; use values lower than 1 to prevent singularities.

sketchkit.image2sketch.swiftsketch.diffusion.gaussian_diffusion.get_named_beta_schedule(schedule_name, num_diffusion_timesteps, scale_betas=1.0, cos_power=2)[source]¶

Get a pre-defined beta schedule for the given name.

The beta schedule library consists of beta schedules which remain similar in the limit of num_diffusion_timesteps. Beta schedules may be added, but should not be removed or changed once they are committed to maintain backwards compatibility.

sketchkit.image2sketch.swiftsketch.diffusion.logger module¶

Logger copied from OpenAI baselines to avoid extra RL-based dependencies: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py

class sketchkit.image2sketch.swiftsketch.diffusion.logger.CSVOutputFormat(filename)[source]¶

Bases: KVWriter

close()[source]¶

writekvs(kvs)[source]¶

class sketchkit.image2sketch.swiftsketch.diffusion.logger.HumanOutputFormat(filename_or_file)[source]¶

Bases: KVWriter, SeqWriter

_truncate(s)[source]¶

close()[source]¶

writekvs(kvs)[source]¶

writeseq(seq)[source]¶

class sketchkit.image2sketch.swiftsketch.diffusion.logger.JSONOutputFormat(filename)[source]¶

Bases: KVWriter

close()[source]¶

writekvs(kvs)[source]¶

class sketchkit.image2sketch.swiftsketch.diffusion.logger.KVWriter[source]¶

Bases: object

writekvs(kvs)[source]¶

class sketchkit.image2sketch.swiftsketch.diffusion.logger.Logger(dir, output_formats, comm=None)[source]¶

Bases: object

CURRENT = None¶

DEFAULT = None¶

_do_log(args)[source]¶

close()[source]¶

dumpkvs()[source]¶

get_dir()[source]¶

log(*args, level=20)[source]¶

logkv(key, val)[source]¶

logkv_mean(key, val)[source]¶

set_comm(comm)[source]¶

set_level(level)[source]¶

class sketchkit.image2sketch.swiftsketch.diffusion.logger.SeqWriter[source]¶

Bases: object

writeseq(seq)[source]¶

class sketchkit.image2sketch.swiftsketch.diffusion.logger.TensorBoardOutputFormat(dir)[source]¶

Bases: KVWriter

Dumps key/value pairs into TensorBoard’s numeric format.

close()[source]¶

writekvs(kvs)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger._configure_default_logger()[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.configure(dir=None, format_strs=None, comm=None, log_suffix='')[source]¶: If comm is provided, average all numerical stats across that comm

sketchkit.image2sketch.swiftsketch.diffusion.logger.debug(*args)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.dump_tabular()¶: Write all of the diagnostics from the current iteration

sketchkit.image2sketch.swiftsketch.diffusion.logger.dumpkvs()[source]¶: Write all of the diagnostics from the current iteration

sketchkit.image2sketch.swiftsketch.diffusion.logger.error(*args)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.get_current()[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.get_dir()[source]¶: Get directory that log files are being written to. will be None if there is no output directory (i.e., if you didn’t call start)

sketchkit.image2sketch.swiftsketch.diffusion.logger.get_rank_without_mpi_import()[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.getkvs()[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.info(*args)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.log(*args, level=20)[source]¶: Write the sequence of args, with no separators, to the console and output files (if you’ve configured an output file).

sketchkit.image2sketch.swiftsketch.diffusion.logger.logkv(key, val)[source]¶: Log a value of some diagnostic Call this once for each diagnostic quantity, each iteration If called many times, last value will be used.

sketchkit.image2sketch.swiftsketch.diffusion.logger.logkv_mean(key, val)[source]¶: The same as logkv(), but if called many times, values averaged.

sketchkit.image2sketch.swiftsketch.diffusion.logger.logkvs(d)[source]¶: Log a dictionary of key-value pairs

sketchkit.image2sketch.swiftsketch.diffusion.logger.make_output_format(format, ev_dir, log_suffix='')[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.mpi_weighted_mean(comm, local_name2valcount)[source]¶: Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110 Perform a weighted average over dicts that are each on a different node Input: local_name2valcount: dict mapping key -> (value, count) Returns: key -> mean

sketchkit.image2sketch.swiftsketch.diffusion.logger.profile(n)[source]¶: Usage: @profile(“my_func”) def my_func(): code

sketchkit.image2sketch.swiftsketch.diffusion.logger.profile_kv(scopename)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.record_tabular(key, val)¶: Log a value of some diagnostic Call this once for each diagnostic quantity, each iteration If called many times, last value will be used.

sketchkit.image2sketch.swiftsketch.diffusion.logger.reset()[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.scoped_configure(dir=None, format_strs=None, comm=None)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.set_comm(comm)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.logger.set_level(level)[source]¶: Set logging threshold on current logger.

sketchkit.image2sketch.swiftsketch.diffusion.logger.warn(*args)[source]¶

sketchkit.image2sketch.swiftsketch.diffusion.resample module¶

class sketchkit.image2sketch.swiftsketch.diffusion.resample.LossAwareSampler[source]¶

Bases: ScheduleSampler

_abc_impl = <_abc._abc_data object>¶

abstractmethod update_with_all_losses(ts, losses)[source]¶

Update the reweighting using losses from a model.

Sub-classes should override this method to update the reweighting using losses from the model.

This method directly updates the reweighting without synchronizing between workers. It is called by update_with_local_losses from all ranks with identical arguments. Thus, it should have deterministic behavior to maintain state across workers.

Parameters:

ts – a list of int timesteps.
losses – a list of float losses, one per timestep.

update_with_local_losses(local_ts, local_losses)[source]¶

Update the reweighting using losses from a model.

Call this method from each rank with a batch of timesteps and the corresponding losses for each of those timesteps. This method will perform synchronization to make sure all of the ranks maintain the exact same reweighting.

Parameters:

local_ts – an integer Tensor of timesteps.
local_losses – a 1D Tensor of losses.

class sketchkit.image2sketch.swiftsketch.diffusion.resample.LossSecondMomentResampler(diffusion, history_per_term=10, uniform_prob=0.001)[source]¶

Bases: LossAwareSampler

_abc_impl = <_abc._abc_data object>¶

_warmed_up()[source]¶

update_with_all_losses(ts, losses)[source]¶

Update the reweighting using losses from a model.

Sub-classes should override this method to update the reweighting using losses from the model.

Parameters:

ts – a list of int timesteps.
losses – a list of float losses, one per timestep.

weights()[source]¶

Get a numpy array of weights, one per diffusion step.

The weights needn’t be normalized, but must be positive.

class sketchkit.image2sketch.swiftsketch.diffusion.resample.ScheduleSampler[source]¶

Bases: ABC

A distribution over timesteps in the diffusion process, intended to reduce variance of the objective.

By default, samplers perform unbiased importance sampling, in which the objective’s mean is unchanged. However, subclasses may override sample() to change how the resampled terms are reweighted, allowing for actual changes in the objective.

_abc_impl = <_abc._abc_data object>¶

sample(batch_size, device)[source]¶

Importance-sample timesteps for a batch.

Parameters:

batch_size – the number of timesteps.
device – the torch device to save to.

Returns:

a tuple (timesteps, weights): - timesteps: a tensor of timestep indices. - weights: a tensor of weights to scale the resulting losses.

abstractmethod weights()[source]¶

Get a numpy array of weights, one per diffusion step.

The weights needn’t be normalized, but must be positive.

class sketchkit.image2sketch.swiftsketch.diffusion.resample.UniformSampler(diffusion)[source]¶

Bases: ScheduleSampler

_abc_impl = <_abc._abc_data object>¶

weights()[source]¶

Get a numpy array of weights, one per diffusion step.

The weights needn’t be normalized, but must be positive.

sketchkit.image2sketch.swiftsketch.diffusion.resample.create_named_schedule_sampler(name, diffusion)[source]¶

Create a ScheduleSampler from a library of pre-defined samplers.

Parameters:

name – the name of the sampler.
diffusion – the diffusion object to sample for.

sketchkit.image2sketch.swiftsketch.diffusion.respace module¶

class sketchkit.image2sketch.swiftsketch.diffusion.respace.SpacedDiffusion(args, use_timesteps, **kwargs)[source]¶

Bases: GaussianDiffusion

A diffusion process which can skip steps in a base diffusion process.

Parameters:

use_timesteps – a collection (sequence or set) of timesteps from the original diffusion process to retain.
kwargs – the kwargs to create the base diffusion process.

_scale_timesteps(t)[source]¶

_wrap_model(model)[source]¶

condition_mean(cond_fn, *args, **kwargs)[source]¶

condition_score(cond_fn, *args, **kwargs)[source]¶

p_mean_variance(model, *args, **kwargs)[source]¶

Apply the model to get p(x_{t-1} | x_t), as well as a prediction of the initial x, x_0.

Parameters:

model – the model, which takes a signal and a batch of timesteps as input.
x – the [N x C x …] tensor at time t.
t – a 1-D Tensor of timesteps.
clip_denoised – if True, clip the denoised signal into [-1, 1].
denoised_fn – if not None, a function which applies to the x_start prediction before it is used to sample. Applies before clip_denoised.
image_features – features used to condition the model’s prediction.

Returns:

plot_forward_pass_with_predicted_x0(model, data, device, scale)[source]¶

training_losses(model, *args, **kwargs)[source]¶

Compute training losses for a single timestep.

Parameters:

model – the model to evaluate loss on.
x_start – the [N x C x …] tensor of inputs. (control points)
x_start_randered_image – the [N x C x …] tensor of inputs_randered_images. (the sketches for image loss)
t – a batch of timestep indices.

Some mean or variance settings may also have other keys.

class sketchkit.image2sketch.swiftsketch.diffusion.respace._WrappedModel(model, timestep_map, rescale_timesteps, original_num_steps)[source]¶: Bases: object

sketchkit.image2sketch.swiftsketch.diffusion.respace.space_timesteps(num_timesteps, section_counts)[source]¶

Create a list of timesteps to use from an original diffusion process, given the number of timesteps we want to take from equally-sized portions of the original process.

For example, if there’s 300 timesteps and the section counts are [10,15,20] then the first 100 timesteps are strided to be 10 timesteps, the second 100 are strided to be 15 timesteps, and the final 100 are strided to be 20.

If the stride is a string starting with “ddim”, then the fixed striding from the DDIM paper is used, and only one section is allowed.

Parameters:

num_timesteps – the number of diffusion steps in the original process to divide up.
section_counts – either a list of numbers, or a string containing comma-separated numbers, indicating the step count per section. As a special case, use “ddimN” where N is a number of steps to use the striding from the DDIM paper.

Returns:

a set of diffusion steps from the original process to use.

Module contents¶

Diffusion utilities used by the SwiftSketch backend.