diffusion_models.models.openai_unetΒΆ

Classes

AttentionBlock(*args, **kwargs)

An attention block that allows spatial positions to attend to each other.

Downsample(*args, **kwargs)

A downsampling layer with an optional convolution.

QKVAttention(*args, **kwargs)

A module which performs QKV attention.

ResBlock(*args, **kwargs)

A residual block that can optionally change the number of channels.

SuperResModel(*args, **kwargs)

A UNetModel that performs super-resolution.

TimestepBlock(*args, **kwargs)

Any module where forward() takes timestep embeddings as a second argument.

TimestepEmbedSequential(*args, **kwargs)

A sequential module that passes timestep embeddings to the children that support it as an extra input.

UNetModel(*args, **kwargs)

The full UNet model with attention and timestep embedding.

Upsample(*args, **kwargs)

An upsampling layer with an optional convolution.