Architectures / Generative
🎨 Diffusion UNet
Stable-Diffusion-style noise predictor — latent UNet with cross-attention to a text embedding
Layers
19
Parameters
6.69M
Input
4 × 64 × 64
Output
4 × 64 × 64
Verifier
Clean
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
19 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 | to_tokens | Reshape | shape=[4096, 320] | 4096 × 320 |
| 7 | down1_text_attn | Cross-Attention | embedDim=320, numHeads=8 | 4096 × 320 |
| 8 | to_feature_map | Reshape | shape=[320, 64, 64] | 320 × 64 × 64 |
| 9 | downsample_1 | Conv2D | outChannels=640, kernelSize=3, stride=2 | 640 × 32 × 32 |
| 10 | mid_norm | GroupNorm | 640 × 32 × 32 | |
| 11 | to_tokens | Reshape | shape=[1024, 640] | 1024 × 640 |
| 12 | mid_text_attn | Cross-Attention | embedDim=640, numHeads=8 | 1024 × 640 |
| 13 | to_feature_map | Reshape | shape=[640, 32, 32] | 640 × 32 × 32 |
| 14 | upsample_1 | Upsample | 640 × 64 × 64 | |
| 15 | up1_conv | Conv2D | outChannels=320, kernelSize=3, stride=1 | 320 × 64 × 64 |
| 16 | conv_out_norm | GroupNorm | 320 × 64 × 64 | |
| 17 | up1_silu | Swish | 320 × 64 × 64 | |
| 18 | conv_out | Conv2D | outChannels=4, kernelSize=3, stride=1 | 4 × 64 × 64 |
| 19 | 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.
No finding. Shapes propagate end to end, every divisibility condition holds, and no advisory rule fires. See the checks.
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.groupNorm_3 = nn.GroupNorm(32, 320)
self.swish_2 = nn.SiLU()
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)
reshape_tokens = swish_1_silu.reshape(swish_1_silu.size(0), 4096, 320)
cross_attention__xattn = self.crossAttention_1(reshape_tokens, reshape_tokens, reshape_tokens)[0]
reshape_to_map = cross_attention__xattn.reshape(cross_attention__xattn.size(0), 320, 64, 64)
conv2d_mple_1 = self.conv2d_3(reshape_to_map)
group_norm_d_norm = self.groupNorm_2(conv2d_mple_1)
reshape_tokens = group_norm_d_norm.reshape(group_norm_d_norm.size(0), 1024, 640)
cross_attention__xattn = self.crossAttention_2(reshape_tokens, reshape_tokens, reshape_tokens)[0]
reshape_to_map = cross_attention__xattn.reshape(cross_attention__xattn.size(0), 640, 32, 32)
upsample_mple_1 = self.upsample_1(reshape_to_map)
conv2d_1_conv = self.conv2d_4(upsample_mple_1)
group_norm_t_norm = self.groupNorm_3(conv2d_1_conv)
swish_1_silu = self.swish_2(group_norm_t_norm)
conv2d_nv_out = self.conv2d_5(swish_1_silu)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.