unet_native

Native UNet for Stable Diffusion

Description

Native R torch implementation of UNet2DConditionModel. Replaces TorchScript for better GPU compatibility.

Usage

unet_native(in_channels = 4L, out_channels = 4L,
            block_out_channels = c(320L, 640L, 1280L, 1280L),
            layers_per_block = 2L, cross_attention_dim = 1024L,
            attention_head_dim = 64L)

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
  • cross_attention_dim: Context dimension from text encoder
  • attention_head_dim: Dimension per attention head

Value

An nn_module representing the UNet