AttentionUNet

class deeplay.components.diffusion.attention_unet.AttentionUNet(*args, **kwargs)

Bases: DeeplayModule

Attention UNet.

Parameters

in_channelsint

Number of input channels.

channelsList[int]

Number of channels in the encoder and decoder blocks.

channel_attentionList[bool]

Attention flags for the encoder and decoder blocks. If True, attention will be applied to the corresponding block. The first attention flag will be ignored as the time information is not integrated at this step. It is still included in the channel_attention just for the sake of consistency.

base_channelsList[int]

Number of channels in the base blocks.

out_channelsint

Number of output channels.

position_embedding_dimint

Dimension of the positional encoding. Positional encoding is defined outside the model and passed as an input to the model. The dimension of the positional encoding should match the dimension given to the model.

num_classesOptional[int]

Number of classes. If num_classes are provided, the class embedding will be added to the positional encoding. This is used for the class conditioned models.

context_embedding_dimOptional[int]

Dimension of the context embedding. Context embedding is defined outside the model and passed as an input to the model. The dimension of the context embedding should match the dimension given to the model. When enabled, the context embedding will be used to apply cross-attention to the feature maps.

num_attention_headsdict

Number of attention heads for self-attention and cross-attention mechanisms. The keys should be “self” and “cross” respectively. Default is {“self”: 1, “cross”: 1}.

Methods Summary

configure(*args, **kwargs)	Configures the module with specified arguments.
forward(x, t[, y, context])	Define the computation performed at every call.

Methods Documentation

configure(*args: Any, **kwargs: Any)

Configures the module with specified arguments.

This method allows dynamic configuration of the module’s properties and behaviors. It can be used to set or modify the attributes and parameters of the module and, if applicable, its child modules. The method intelligently handles both direct attribute configuration and delegation to child modules’ configure methods.

Parameters

*argsAny

Positional arguments specifying the configuration settings. When the first argument is a string matching a configurable attribute, the method expects either one or two arguments: the attribute name and, optionally, its value. If the attribute is itself a DeeplayModule, subsequent arguments are passed to its configure method.

**kwargsAny

Keyword arguments for configuration settings. If provided, these are used to update the module’s configuration directly.

Raises

ValueError

Raised if a configuration key is not recognized as a valid configurable for the module or if the provided arguments do not match the expected pattern for configuration.

Example Usage

To configure a single attribute: ` module.configure('attribute_name', attribute_value) # or module.configure(attribute_name=attribute_value) `

To configure multiple attributes using keyword arguments: ` module.configure(attribute1=value1, attribute2=value2) `

To configure a child module’s attribute: ` module.configure('child_module_attribute', child_attribute=child_attribute_value) # or module.child_module.configure(child_attribute=child_attribute_value) `

forward(x, t, y=None, context=None)

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.