diffusion_models.models.openai_unet.AttentionBlock

class diffusion_models.models.openai_unet.AttentionBlock(*args: Any, **kwargs: Any)[source]

Bases: Module

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

Originally ported from here, but adapted to the N-d case. https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.

__init__(channels, num_heads=1, use_checkpoint=False)[source]

Methods

__init__(channels[, num_heads, use_checkpoint])

forward(x)