# Diffusion UNet

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

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.

- Category: Generative
- Layers: 19
- Parameters: 6.69M
- Input shape (batchless): 4 × 64 × 64
- Output shape: 4 × 64 × 64
- Verifier verdict: pass
- Graph JSON: https://neurarch.com/templates/diffusion-unet/model.json
- Open on the canvas: https://neurarch.com/?template=diffusion-unet

## Structure

| # | 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 |

## Verifier findings

No finding. Shapes propagate end to end and no advisory rule fires.

## Exported PyTorch (first 46 lines)

```python
# 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)
```

## Machine access

- Every architecture: https://neurarch.com/a/index.json
- Verify a graph of your own: `POST https://www.neurarch.com/api/v1/check` (see https://neurarch.com/developer.html)
- MCP server, so an agent edits the graph with the checks in the loop: https://neurarch.com/docs/mcp.md