nx_unet_block.nx

nx_unet_block.nx -- diffusion-model ResBlock composer. L4 composer for the UNet building block used by Stable Diffusion 1.x/2.x/3, Flux, Z-Image. Pure composition -- no new math. ===== Block shape (canonical SD ResBlock) ======================= x: [N, C, H, W] input 1. h = GroupNorm(x, num_groups=32, gamma_1, beta_1) 2. h = SiLU(h) 3. h = Conv2D(h, W_1) : [N, C, H, W] 4. h = GroupNorm(h, num_groups=32, gamma_2, beta_2) 5. h = SiLU(h) 6. h = Conv2D(h, W_2) : [N, C, H, W] 7. Skip: residual = x (v1: assumes C_in = C_out) 8. Output = h + residual ===== v1 scope ================================================== Assumes C_in == C_out throughout (skip-connection is identity). For "down" blocks in a U-shape (C_in != C_out): caller passes a separate 1x1 skip-conv layer and we'd extend this composer in v2. v1 satisfies the same-channel-count case, which is the majority of ResBlocks within a stage. Time-embedding injection (diffusion-step conditioning): canonical UNets add a per-block timestep-embedding vector to h between steps 3 and 4 (after first conv, before second norm). v1 omits this -- caller is expected to add it post-hoc or use the queued v2 of this primitive that accepts a time_emb tensor. Bits-up composition (pure, no new math): NxTensor (L1) nx_conv2d (L3) nx_groupnorm (L3, shipped 40d6e1c2) nx_silu (L3) nx_loop.LoopVerdict (control) genealogy_id: ho_2020_ddpm_unet + ronneberger_2015_u_net +