Nestmodels for Linnet
Models / sdxl-base-unet

SDXL Base UNet

The 2.6B-parameter denoising UNet of Stable Diffusion XL base: residual blocks and spatial transformers at three resolutions, cross-attending to a 2048-wide text embedding, conditioned on the timestep and on SDXL's pooled-text-plus-micro-conditioning vector.

2.6B parameterssdxlCreativeML Open RAIL++-Mimage-generationvisionconvolutionaldiffusiontext-to-imageunet
Hugging FaceGitHubarXiv Registry entry

The forward entry with one level of blocks expanded. Every edge carries the tensor type the compiler inferred at that point, in the model's own generics.

SDXL Base UNet: forwardSDXL Base UNet: forward

Entries ​

EntrySignature
forwardforward<B: Dim, S: Dim>(latent: Tensor[B, 4, 4 * MidH, 4 * MidW; T], timestep_embedding: Tensor[B, 320; T], text_embedding: Tensor[B, S, 2048; T], added_embedding: Tensor[B, 2816; T]) -> Tensor[B, 4, 4 * MidH, 4 * MidW; T]

Generics ​

The root block's generics as this checkpoint binds them.

MidH8
MidW8
Depth12
Depth210
Tf32

Blocks ​

Every block of the program with its members and functions, as linnet inspect prints them.

Attention ​

text
sdxl_unet::Attention<C: Dim, Ctxdim: Dim, Heads: Dim, T: Float>
  sub to_q: Linear<C, C, T>
  sub to_k: Linear<Ctxdim, C, T>
  sub to_v: Linear<Ctxdim, C, T>
  sub to_out: Linear<C, C, T>
  pub fn forward<B: Dim, N: Dim, S: Dim>(x: Tensor[B, N, C; T], context: Tensor[B, S, Ctxdim; T]) -> Tensor[B, N, C; T]

Conv1d ​

text
std.nn.conv::Conv1d<Cin: Dim, Cout: Dim, K: Dim, Stride: Dim, Pad: Dim, T: Float = f32>
  param weight: Tensor[Cout, Cin, K; T]
  param bias: Tensor[Cout; T]?
  pub fn forward<B: Dim, L: Dim>(x: Tensor[B, Cin, L; T]) -> Tensor[B, Cout, 1 + (-1 * K + L + 2 * Pad) / Stride; T]

Conv2d ​

text
std.nn.conv::Conv2d<Cin: Dim, Cout: Dim, K: Dim, Stride: Dim, Pad: Dim, T: Float = f32>
  param weight: Tensor[Cout, Cin, K, K; T]
  param bias: Tensor[Cout; T]?
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, Cin, H, W; T]) -> Tensor[B, Cout, 1 + (-1 * K + H + 2 * Pad) / Stride, 1 + (-1 * K + W + 2 * Pad) / Stride; T]

Conv2dRect ​

text
std.nn.conv::Conv2dRect<Cin: Dim, Cout: Dim, KH: Dim, KW: Dim, StrideH: Dim, StrideW: Dim, PadH: Dim, PadW: Dim, T: Float = f32>
  param weight: Tensor[Cout, Cin, KH, KW; T]
  param bias: Tensor[Cout; T]?
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, Cin, H, W; T]) -> Tensor[B, Cout, 1 + (-1 * KH + H + 2 * PadH) / StrideH, 1 + (-1 * KW + W + 2 * PadW) / StrideW; T]

FeedForward ​

text
sdxl_unet::FeedForward<C: Dim, Inner: Dim, T: Float>
  sub proj: Linear<C, 2 * Inner, T>
  sub out: Linear<Inner, C, T>
  pub fn forward<B: Dim, N: Dim>(x: Tensor[B, N, C; T]) -> Tensor[B, N, C; T]

GroupNorm ​

text
sdxl_unet::GroupNorm<C: Dim, T: Float>
  param weight: Tensor[C; T]
  param bias: Tensor[C; T]
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, C, H, W; T], eps: f32) -> Tensor[B, C, H, W; T]

LayerNorm ​

text
sdxl_unet::LayerNorm<C: Dim, T: Float>
  param weight: Tensor[C; T]
  param bias: Tensor[C; T]
  pub fn forward<*S: Shape>(x: Tensor[*S, C; T]) -> Tensor[*S, C; T]

Linear ​

text
std.nn.linear::Linear<In: Dim, Out: Dim, T: Float = bf16>
  param weight: Tensor[Out, In; T]
  param bias: Tensor[Out; T]?
  pub fn forward<*S: Shape>(x: Tensor[*S, In; T]) -> Tensor[*S, Out; T]

Model ​

text
sdxl_unet::Model<MidH: Dim, MidW: Dim, Depth1: Dim, Depth2: Dim, T: Float = f32>
  sub conv_in: Conv2d<4, 320, 3, 1, 1, T>
  sub time_linear_1: Linear<320, 1280, T>
  sub time_linear_2: Linear<1280, 1280, T>
  sub add_linear_1: Linear<2816, 1280, T>
  sub add_linear_2: Linear<1280, 1280, T>
  sub down0_res0: ResBlock<320, 1280, T>
  sub down0_res1: ResBlock<320, 1280, T>
  sub down0_downsample: Conv2d<320, 320, 3, 2, 1, T>
  sub down1_res0: ResBlockWiden<320, 640, 1280, T>
  sub down1_attn0: Transformer<640, 2048, 10, Depth1, T>
  sub down1_res1: ResBlock<640, 1280, T>
  sub down1_attn1: Transformer<640, 2048, 10, Depth1, T>
  sub down1_downsample: Conv2d<640, 640, 3, 2, 1, T>
  sub down2_res0: ResBlockWiden<640, 1280, 1280, T>
  sub down2_attn0: Transformer<1280, 2048, 20, Depth2, T>
  sub down2_res1: ResBlock<1280, 1280, T>
  sub down2_attn1: Transformer<1280, 2048, 20, Depth2, T>
  sub mid_res0: ResBlock<1280, 1280, T>
  sub mid_attn: Transformer<1280, 2048, 20, Depth2, T>
  sub mid_res1: ResBlock<1280, 1280, T>
  sub up0_res0: ResBlockWiden<2560, 1280, 1280, T>
  sub up0_attn0: Transformer<1280, 2048, 20, Depth2, T>
  sub up0_res1: ResBlockWiden<2560, 1280, 1280, T>
  sub up0_attn1: Transformer<1280, 2048, 20, Depth2, T>
  sub up0_res2: ResBlockWiden<1920, 1280, 1280, T>
  sub up0_attn2: Transformer<1280, 2048, 20, Depth2, T>
  sub up0_upsample: Conv2d<1280, 1280, 3, 1, 1, T>
  sub up1_res0: ResBlockWiden<1920, 640, 1280, T>
  sub up1_attn0: Transformer<640, 2048, 10, Depth1, T>
  sub up1_res1: ResBlockWiden<1280, 640, 1280, T>
  sub up1_attn1: Transformer<640, 2048, 10, Depth1, T>
  sub up1_res2: ResBlockWiden<960, 640, 1280, T>
  sub up1_attn2: Transformer<640, 2048, 10, Depth1, T>
  sub up1_upsample: Conv2d<640, 640, 3, 1, 1, T>
  sub up2_res0: ResBlockWiden<960, 320, 1280, T>
  sub up2_res1: ResBlockWiden<640, 320, 1280, T>
  sub up2_res2: ResBlockWiden<640, 320, 1280, T>
  sub norm_out: GroupNorm<320, T>
  sub conv_out: Conv2d<320, 4, 3, 1, 1, T>
  pub entry forward<B: Dim, S: Dim>(latent: Tensor[B, 4, 4 * MidH, 4 * MidW; T], timestep_embedding: Tensor[B, 320; T], text_embedding: Tensor[B, S, 2048; T], added_embedding: Tensor[B, 2816; T]) -> Tensor[B, 4, 4 * MidH, 4 * MidW; T]

ResBlock ​

text
sdxl_unet::ResBlock<C: Dim, Temb: Dim, T: Float>
  sub norm1: GroupNorm<C, T>
  sub conv1: Conv2d<C, C, 3, 1, 1, T>
  sub time_emb_proj: Linear<Temb, C, T>
  sub norm2: GroupNorm<C, T>
  sub conv2: Conv2d<C, C, 3, 1, 1, T>
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, C, H, W; T], conditioning: Tensor[B, Temb; T]) -> Tensor[B, C, H, W; T]

ResBlockWiden ​

text
sdxl_unet::ResBlockWiden<Cin: Dim, Cout: Dim, Temb: Dim, T: Float>
  sub norm1: GroupNorm<Cin, T>
  sub conv1: Conv2d<Cin, Cout, 3, 1, 1, T>
  sub time_emb_proj: Linear<Temb, Cout, T>
  sub norm2: GroupNorm<Cout, T>
  sub conv2: Conv2d<Cout, Cout, 3, 1, 1, T>
  sub shortcut: Conv2d<Cin, Cout, 1, 1, 0, T>
  pub fn forward<B: Dim, H: Dim, W: Dim>(x: Tensor[B, Cin, H, W; T], conditioning: Tensor[B, Temb; T]) -> Tensor[B, Cout, H, W; T]

RmsNorm ​

text
std.nn.norm::RmsNorm<H: Dim, T: Float = bf16>
  param weight: Tensor[H; T]
  pub fn forward<*S: Shape>(x: Tensor[*S, H; T]) -> Tensor[*S, H; T]

Transformer ​

text
sdxl_unet::Transformer<C: Dim, Ctxdim: Dim, Heads: Dim, Depth: Dim, T: Float>
  sub norm: GroupNorm<C, T>
  sub proj_in: Linear<C, C, T>
  sub blocks: [TransformerBlock<C, Ctxdim, Heads, T>; Depth]
  sub proj_out: Linear<C, C, T>
  pub fn forward<B: Dim, H: Dim, W: Dim, S: Dim>(x: Tensor[B, C, H, W; T], context: Tensor[B, S, Ctxdim; T]) -> Tensor[B, C, H, W; T]

TransformerBlock ​

text
sdxl_unet::TransformerBlock<C: Dim, Ctxdim: Dim, Heads: Dim, T: Float>
  sub norm1: LayerNorm<C, T>
  sub attn1: Attention<C, C, Heads, T>
  sub norm2: LayerNorm<C, T>
  sub attn2: Attention<C, Ctxdim, Heads, T>
  sub norm3: LayerNorm<C, T>
  sub ff: FeedForward<C, 4 * C, T>
  pub fn forward<B: Dim, N: Dim, S: Dim>(x: Tensor[B, N, C; T], context: Tensor[B, S, Ctxdim; T]) -> Tensor[B, N, C; T]