unet_sdxl_native

Packages

Native SDXL UNet

Description

Native R torch implementation of SDXL UNet2DConditionModel. SDXL has a different architecture from SD21:

3 down/up blocks (not 4)
Variable transformer depth per block
Additional conditioning via add_embedding

Usage

unet_sdxl_native(in_channels = 4L, out_channels = 4L,
                 block_out_channels = c(320L, 640L, 1280L),
                 layers_per_block = 2L,
                 transformer_layers_per_block = c(0L, 2L, 10L),
                 cross_attention_dim = 2048L, attention_head_dim = 64L,
                 addition_embed_dim = 1280L, addition_time_embed_dim = 256L)

Arguments

in_channels: Input channels (default 4 for latent space)
out_channels: Output channels (default 4)
block_out_channels: Channel multipliers per block
layers_per_block: Number of ResBlocks per down/up block
transformer_layers_per_block: Transformer depth per block
cross_attention_dim: Context dimension from text encoder
attention_head_dim: Dimension per attention head
addition_embed_dim: Dimension for additional embeddings
addition_time_embed_dim: Dimension for time embedding projection

Value

An nn_module representing the SDXL UNet