Architectures / Generative

🎨 Diffusion UNet

Stable-Diffusion-style noise predictor — latent UNet with cross-attention to a text embedding

Layers

Parameters

6.69M

Input

4 × 64 × 64

Output

4 × 64 × 64

Verifier

3 advisories

Every number on this page is computed from the graph by the same functions the app runs, not written by hand.

Open Diffusion UNet on the canvas Free, no account needed

When to pick it

Pick when you want to generate images from a text prompt. The full pipeline also needs a VAE encoder/decoder and a text encoder (e.g. CLIP); this template is the denoiser core.

Structure

15 layers. Output shapes are propagated from the input shape, batch dimension excluded.

	Layer	Type	Parameters	Output shape
1	noisy_latent	Input	shape=[4, 64, 64]	4 × 64 × 64
2	conv_in	Conv2D	outChannels=320, kernelSize=3, stride=1	320 × 64 × 64
3	down1_norm	GroupNorm		320 × 64 × 64
4	down1_conv	Conv2D	outChannels=320, kernelSize=3, stride=1	320 × 64 × 64
5	down1_silu	Swish		320 × 64 × 64
6	down1_text_attn	Cross-Attention	embedDim=320, numHeads=8	320 × 64 × 64
7	downsample_1	Conv2D	outChannels=640, kernelSize=3, stride=2	640 × 32 × 32
8	mid_norm	GroupNorm		640 × 32 × 32
9	mid_text_attn	Cross-Attention	embedDim=640, numHeads=8	640 × 32 × 32
10	upsample_1	Upsample		640 × 64 × 64
11	up1_conv	Conv2D	outChannels=320, kernelSize=3, stride=1	320 × 64 × 64
12	up1_silu	Swish		320 × 64 × 64
13	conv_out_norm	GroupNorm		320 × 64 × 64
14	conv_out	Conv2D	outChannels=4, kernelSize=3, stride=1	4 × 64 × 64
15	predicted_noise	Output		4 × 64 × 64

What the verifier says

The same 41 structural checks that run on every edit in the app, on this graph.

warn"conv_out_norm" (groupNorm) follows "up1_silu" (activation). The standard pre-activation order is Conv/Linear → Norm → Activation. Normalizing post-activation limits expressivity. Fix: Move the normalization layer before the activation function. (up1_silu)
bn-after-activation

warncrossAttention 'down1_text_attn': embedDim=320 but upstream last dim is 64; attention expects them equal (project the input or fix embedDim) (down1_text_attn)
attention-in-mismatch

warncrossAttention 'mid_text_attn': embedDim=640 but upstream last dim is 32; attention expects them equal (project the input or fix embedDim) (mid_text_attn)
attention-in-mismatch

The PyTorch it exports

Generated from the graph above. First 46 lines; the app exports the whole file, plus the training loop, the data contract and a deploy bundle.

# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class DiffusionUNetStable_Diffusion_style(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv2d_1 = nn.Conv2d(4, 320, kernel_size=3, stride=1, padding=1)
        self.groupNorm_1 = nn.GroupNorm(32, 320)
        self.conv2d_2 = nn.Conv2d(320, 320, kernel_size=3, stride=1, padding=1)
        self.swish_1 = nn.SiLU()
        self.crossAttention_1 = nn.MultiheadAttention(embed_dim=320, num_heads=8, batch_first=True)
        self.conv2d_3 = nn.Conv2d(320, 640, kernel_size=3, stride=2, padding=1)
        self.groupNorm_2 = nn.GroupNorm(32, 640)
        self.crossAttention_2 = nn.MultiheadAttention(embed_dim=640, num_heads=8, batch_first=True)
        self.upsample_1 = nn.Upsample(scale_factor=2, mode='nearest')
        self.conv2d_4 = nn.Conv2d(640, 320, kernel_size=3, stride=1, padding=1)
        self.swish_2 = nn.SiLU()
        self.groupNorm_3 = nn.GroupNorm(32, 320)
        self.conv2d_5 = nn.Conv2d(320, 4, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        # noisy_latent shape: [4,64,64]
        conv2d_onv_in = self.conv2d_1(x)
        group_norm_1_norm = self.groupNorm_1(conv2d_onv_in)
        conv2d_1_conv = self.conv2d_2(group_norm_1_norm)
        swish_1_silu = self.swish_1(conv2d_1_conv)
        cross_attention__xattn = self.crossAttention_1(swish_1_silu, swish_1_silu, swish_1_silu)[0]
        conv2d_mple_1 = self.conv2d_3(cross_attention__xattn)
        group_norm_d_norm = self.groupNorm_2(conv2d_mple_1)
        cross_attention__xattn = self.crossAttention_2(group_norm_d_norm, group_norm_d_norm, group_norm_d_norm)[0]
        upsample_mple_1 = self.upsample_1(cross_attention__xattn)
        conv2d_1_conv = self.conv2d_4(upsample_mple_1)
        swish_1_silu = self.swish_2(conv2d_1_conv)
        group_norm_t_norm = self.groupNorm_3(swish_1_silu)
        conv2d_nv_out = self.conv2d_5(group_norm_t_norm)
        # Output
        return conv2d_nv_out

For agents

This architecture is machine-readable end to end. An agent can list the set, fetch this graph, edit it, and have the edit verified before any GPU time is spent.

Markdown twin Index JSON This graph as JSON MCP server POST /api/v1/check

Also in Generative

🌀 DiT-XL/2

Diffusion Transformer — replaces the UNet denoiser with a ViT backbone conditioned on timestep + class via adaLN-Zero

204 layers · 670.68M