File size: 5,184 Bytes
19ff442 861810c 19ff442 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | ---
datasets:
- tiny_imagenet
language: en
library_name: jax
license: mit
pipeline_tag: text-to-image
tags:
- diffusion-model
- flax
- flax-nnx
- jax
- unet
- tiny-imagenet
- conditional-image-generation
- from-scratch
---
# Conditional MediumUNet Diffusion Model (Flax NNX, JAX)
A class-conditional, from-scratch DDPM-style diffusion model trained on **Tiny ImageNet ($64 imes 64 imes 3$, 200 classes)**, built using **JAX** and **Flax NNX**.
The model architecture (`ConditionalUNet`) predicts the clean image from a noisy input, conditioned simultaneously on a sinusoidally embedded continuous noise level and a class label embedding.
## Model Details
- **Architecture:** Medium U-Net with Class Conditioning (featuring sinusoidal time embeddings, learnable class embeddings, strided convolutions for downsampling, transposed convolutions for upsampling, and functional residual skip connections)
- **Conditioning:**
1. Scalar noise level `amount ∈ [0, 1]` via sinusoidal positional embeddings.
2. Class index `label ∈ [0, 199]` via dense embedding vectors injected into the residual blocks.
- **Framework:** JAX / Flax NNX
- **Training data:** Tiny ImageNet (200 classes) — full 100,000-image Tiny ImageNet training set
- **Optimizer:** AdamW with warmup-cosine-decay LR schedule, gradient clipping (global norm 1.0)
- **Epochs:** 120
- **Batch size:** 256
## 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` |
| `conditional_medium_unet.mlir` | MLIR | StableHLO dialect, via `jax.export` |
| `conditional_medium_unet.stablehlo` | StableHLO bytecode | Portable compiled representation (**with 100% baked-in weights**) |
## Usage (JAX / Flax NNX Native)
```python
import orbax.checkpoint as ocp
from flax import nnx
# Recreate the model skeleton, then restore the trained parameters
model = ConditionalUNet(in_features=3, out_features=3, num_classes=200, 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 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 with the **HQ quadratic schedule (80 steps)**:
```python
import jax
import jax.numpy as jnp
import numpy as np
# 1. Load the pre-compiled StableHLO bytecode from file
with open("conditional_medium_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 normal noise (batch_size=1, 64x64, 3 channels)
key = jax.random.PRNGKey(42)
x = jax.random.normal(key, shape=(1, 64, 64, 3))
# Choose class label (e.g., 42 for 'mantis')
CHOSEN_CLASS = 42
sample_label = jnp.array([CHOSEN_CLASS], dtype=jnp.int32)
# 4. Set up the HQ Quadratic Schedule (80 steps)
n_steps = 80
steps_line = jnp.linspace(1.0, 0.0, n_steps + 1)
ts = steps_line ** 2
print("🚀 Running conditional inference loop via independent StableHLO runtime...")
for i in range(n_steps):
t_curr = ts[i]
t_next = ts[i + 1]
current_amount = jnp.array([t_curr], dtype=jnp.float32)
# 🔥 LIGHTNING-FAST HARDWARE INFERENCE:
# Pass image, time step, and class label. No model state or weights needed!
pred_clean = compiled_fn(x, current_amount, sample_label)
pred_clean = jnp.clip(pred_clean, -1.0, 1.0)
if i == n_steps - 1:
x = pred_clean
break
# DDIM Mathematical Correction
eps_estimated = (x - (1.0 - t_curr) * pred_clean) / (t_curr + 1e-8)
x = (1.0 - t_next) * pred_clean + t_next * eps_estimated
# Rescale from [-1, 1] back to [0, 1] for image display
final_image = np.array((x + 1.0) / 2.0)[0]
print("✅ Image generated successfully! Output shape:", final_image.shape)
```
## Technical Parameters & Design Choices
- **Color & Resolution:** Native support for $64 imes 64$ RGB images, a massive step up in spatial complexity compared to standard MNIST benchmarks.
- **Class-Conditional Latent Guidance:** Embedding-driven injection allows targeted generation of specific objects among 200 distinct classes, turning the network from an unconditional sampler into a controllable asset.
- **Production Readiness:** Ideal for specialized target applications on low-power edge nodes or cross-language backends (C++, Rust) via the native StableHLO runtime interface.
|