NanoUNet Diffusion Model (Flax NNX, JAX)
A tiny, from-scratch DDPM-style diffusion model trained on MNIST, built with JAX and Flax NNX.
The model is an optimized conditional U-Net (RealNanoUNet) that predicts the clean image from a noisy
input, conditioned on a sinusoidally embedded noise level.
Model Details
- Architecture: 2-level Enhanced U-Net (32 to 64 features, double convolutions per stage, strided-conv downsampling, transposed-conv upsampling, and functional residual skip connections)
- Conditioning: scalar noise level
amount ∈ [0, 1], embedded with sinusoidal positional embeddings - Framework: JAX / Flax NNX
- Training data: MNIST — {training_data_desc}
- Optimizer: AdamW with warmup-cosine-decay LR schedule, gradient clipping (global norm 1.0)
- Epochs: 100
- Batch size: 128
Exported Formats
This repository bundles multiple export formats of the same trained weights:
| File / Folder | Format | Notes |
|---|---|---|
model_checkpoint/ |
Orbax checkpoint | Native JAX/Flax NNX state, load with orbax.checkpoint |
nano_unet.mlir |
MLIR | StableHLO dialect, via jax.export |
nano_unet.stablehlo |
StableHLO bytecode | Portable compiled representation (with 100% baked-in weights) |
Usage (JAX / Flax NNX Native)
import orbax.checkpoint as ocp
from flax import nnx
# Recreate the model skeleton, then restore the trained parameters
model = RealNanoUNet(in_features=1, out_features=1, rngs=nnx.Rngs(0))
graphdef, abstract_state = nnx.split(model)
checkpointer = ocp.Checkpointer(ocp.StandardCheckpointHandler())
state = checkpointer.restore("model_checkpoint", abstract_state)
model = nnx.merge(graphdef, state)
⚡ How to use the StableHLO Model (Production Inference)
The ultimate advantage of this pipeline is that you don't need the RealNanoUNet Python architecture class anymore, nor do you need to manage raw weight parameters! Because the model was exported via a functional static JIT trace, all trained weights are fully baked directly into the StableHLO bytecode binary. The compiled graph acts as a completely autonomous black-box function.
Here is how to properly load and execute the pre-compiled StableHLO graph to perform a full iterative diffusion reverse loop using pure JAX:
import jax
import jax.numpy as jnp
import numpy as np
# 1. Load the pre-compiled StableHLO bytecode from file
with open("nano_unet.stablehlo", "rb") as f:
stablehlo_bytecode = f.read()
# 2. Deserialize it back into a callable JAX artifact and JIT compile it
loaded_artifact = jax.export.deserialize(stablehlo_bytecode)
compiled_fn = jax.jit(loaded_artifact.call)
# 3. Initialize pure uniform noise (as a starting point for generation)
key = jax.random.PRNGKey(42)
x = jax.random.uniform(key, shape=(1, 28, 28, 1))
n_steps = 20
print("🚀 Running iterative diffusion loop via independent StableHLO runtime...")
# 4. Multi-step reverse sampling process
for i in range(n_steps):
t = 1.0 - (i / n_steps)
current_amount = jnp.array([t], dtype=jnp.float32)
# 🔥 LIGHTNING-FAST HARDWARE INFERENCE:
# Just pass the inputs! No model state or weights needed in the arguments!
pred_clean = compiled_fn(x, current_amount)
pred_clean = jnp.clip(pred_clean, 0.0, 1.0)
if i == n_steps - 1:
x = pred_clean
break
t_next = 1.0 - ((i + 1) / n_steps)
eps_estimated = (x - (1.0 - t) * pred_clean) / (t + 1e-8)
x = (1.0 - t_next) * pred_clean + t_next * eps_estimated
print("✅ Image generated successfully! Output shape:", x.shape)
Technical Parameters & Design Choices
- Architecture: A optimized, double-convoluted U-Net tailored for efficient edge inference and educational deployment. The enhanced capacity prevents geometric pattern artifacts (grid noise) during sampling.
- Noise Process: Uses continuous linear interpolation between data and noise space for faster convergence and streamlined math, making it highly transparent for training.
- Scale: The network scale and hyperparameter configuration are intentionally optimized to allow rapid training and instant multi-platform deployment (StableHLO/MLIR) without requiring enterprise-grade cluster resources.