- added support for google/ddpm-cifar10-32

- click image to show it in its original resolution

.gitattributes

@@ -1 +1,2 @@
 assets/celeba_hq_validation_91_256.png filter=lfs diff=lfs merge=lfs -text
+assets/cifar10_train_23_32.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -11,7 +11,7 @@ pinned: false
 
 # Diffusion Lab
 
-The Gradio app supports ZeroGPU and streams checkpoint images. The reverse process runs the pretrained `google/ddpm-celebahq-256` model through Hugging Face Diffusers. This is an unconditional CelebA-HQ face generator: reverse sampling normally starts with random Gaussian noise. You can also select the final forward-noise image as its input; this is a demonstration, not a guaranteed reconstruction of the uploaded image.
+The Gradio app supports ZeroGPU and streams checkpoint images. The model selector offers `google/ddpm-celebahq-256` (256×256 faces, the default) and `google/ddpm-cifar10-32` (32×32 CIFAR-10 objects and animals) through Hugging Face Diffusers. Forward images and noise use the selected model’s resolution in Gradio. This is an unconditional CelebA-HQ face generator: reverse sampling normally starts with random Gaussian noise. You can also select the final forward-noise image as its input; this is a demonstration, not a guaranteed reconstruction of the uploaded image.
 
 ## Run
 
@@ -24,7 +24,7 @@ python -m pip install -r requirements.txt
 python app.py
 ```
 
-Open <http://127.0.0.1:7860>. Startup downloads the model weights into `models/`; later runs use that cache. The backend prefers CUDA, then Apple Silicon MPS, then CPU. Model inference uses all 1,000 DDPM steps and sends the selected intermediate images to the browser as they are produced.
+Open <http://127.0.0.1:7860>. On first use, each selected model downloads its weights into `models/`; later runs use that cache. The backend prefers CUDA, then Apple Silicon MPS, then CPU. Model inference uses all 1,000 DDPM steps and sends the selected intermediate images to the browser as they are produced.
 
 To generate and display one image without starting the server, run:
 
@@ -40,11 +40,18 @@ python download_celeba_hq.py 1374
 python download_celeba_hq.py 42 --split train --output face-42.png
 ```
 
+To download a 32×32 CIFAR-10 image (test split by default):
+
+```sh
+python download_cifar10.py 1374
+python download_cifar10.py 42 --split train --output assets/cifar10-train-42.png
+```
+
 ## ZeroGPU deployment
 
 Push this repository to a Gradio Space and select ZeroGPU in the Space hardware settings. The Dockerfile and legacy server are not used by the Gradio deployment. Account eligibility and daily GPU quotas apply.
 
-The app imports `spaces` before PyTorch and places the GPU model on CUDA at startup when `SPACES_ZERO_GPU` is enabled. Reverse sampling tries GPU first and automatically restarts on CPU if GPU allocation, quota, or supported CUDA availability errors prevent the run. Locally, the app prefers CUDA, then Apple Silicon MPS, then CPU. MPS availability failures also fall back to CPU. Separate pipelines are cached for each device; the CPU pipeline loads on first use. CPU sampling is substantially slower. GPU sampling runs inside a decorated function called by the Gradio handler. `GPU_DURATION_SECONDS` defaults to 120; benchmark a full 1,000-step run and adjust this Space variable if necessary. A timeout is reported by the Gradio queue; selecting fewer checkpoints does not reduce the sampling steps.
+The app imports `spaces` before PyTorch and places both GPU models on CUDA at startup when `SPACES_ZERO_GPU` is enabled. Reverse sampling tries GPU first and automatically restarts on CPU if GPU allocation, quota, or supported CUDA availability errors prevent the run. Locally, the app prefers CUDA, then Apple Silicon MPS, then CPU. MPS availability failures also fall back to CPU. Separate pipelines are cached for each model and device; the CPU pipeline loads on first use. CPU sampling is substantially slower. GPU sampling runs inside a decorated function called by the Gradio handler. `GPU_DURATION_SECONDS` defaults to 120; benchmark a full 1,000-step run and adjust this Space variable if necessary. A timeout is reported by the Gradio queue; selecting fewer checkpoints does not reduce the sampling steps.
 
 For persistent caching, attach a writable storage bucket and set the Space variable `MODEL_CACHE_DIR` to its mount path (for example `/data/models`). The old Docker path `/app/models` is also usable if explicitly configured with that variable.
 

app.js

@@ -1,5 +1,5 @@
 const TOTAL_STEPS = 1000;
-const MODEL_NAME = 'google/ddpm-celebahq-256';
+const modelInput = document.querySelector('#model-input');
 const DEFAULT_CHECKPOINTS = [0, 50, 100, 200, 350, 500, 700, 1000];
 const SIZE = 192;
 const PIXEL_COUNT = SIZE * SIZE * 3;
@@ -58,7 +58,7 @@ function computeForwardStates() {
   forwardStates = checkpoints.map((step) => { if (step === 0) return new Float32Array(sourcePixels); const alphaBar = schedule[step - 1]; const state = new Float32Array(PIXEL_COUNT); for (let index = 0; index < PIXEL_COUNT; index += 1) state[index] = Math.sqrt(alphaBar) * sourcePixels[index] + Math.sqrt(1 - alphaBar) * noise[index]; return state; });
 }
 function makeNoiseState(seed) { const random = randomNormal(seed); const state = new Float32Array(PIXEL_COUNT); for (let index = 0; index < PIXEL_COUNT; index += 1) state[index] = random(); return state; }
-function stateToPng(state) { const canvas = document.createElement('canvas'); canvas.width = 256; canvas.height = 256; const small = document.createElement('canvas'); small.width = SIZE; small.height = SIZE; drawState(small, state); canvas.getContext('2d').drawImage(small, 0, 0, 256, 256); return canvas.toDataURL('image/png'); }
+function stateToPng(state) { const canvas = document.createElement('canvas'); const size = modelInput.value === 'google/ddpm-cifar10-32' ? 32 : 256; canvas.width = size; canvas.height = size; const small = document.createElement('canvas'); small.width = SIZE; small.height = SIZE; drawState(small, state); canvas.getContext('2d').drawImage(small, 0, 0, size, size); return canvas.toDataURL('image/png'); }
 function setModelStatus(message) { if (elements.modelStatus) elements.modelStatus.textContent = message; }
 function resampleNoise() {
   const values = new Uint32Array(1); crypto.getRandomValues(values); const seed = values[0] & 0x7fffffff; elements.seedInput.value = String(seed); reverseInputMode = 'seed'; reverseNoiseState = makeNoiseState(seed); reverseStates = new Map(); resetProgress('reverse'); setModelStatus(`Resampled Gaussian noise · seed ${seed}.`); renderTracks();
@@ -81,12 +81,29 @@ function makeCard(step, state, type, index) {
   if (type === 'forward' && index === 0) { const change = document.createElement('button'); change.className = 'change-image'; change.type = 'button'; change.textContent = 'Change image'; change.addEventListener('click', () => elements.imageUpload.click()); canvasWrap.append(change); }
   if (type === 'reverse' && index === checkpoints.length - 1) { const noiseActions = document.createElement('div'); noiseActions.className = 'noise-actions'; noiseActions.append(makeActionButton('Resample', 'resample-button', resampleNoise), makeActionButton('Forward noise', 'forward-noise-button', useForwardNoise)); canvasWrap.append(noiseActions); }
   const caption = document.createElement('div'); caption.className = 'stage-caption';
-  const titleText = type === 'forward' && index === 0 ? imageLabel : step === 0 ? 'Generated face' : step === TOTAL_STEPS ? 'Gaussian noise' : '';
+  const titleText = type === 'forward' && index === 0 ? imageLabel : step === 0 ? 'Generated image' : step === TOTAL_STEPS ? 'Gaussian noise' : '';
   if (titleText) { const title = document.createElement('strong'); title.textContent = titleText; caption.append(title); }
   const time = document.createElement('span'); time.textContent = `t = ${step}`; caption.append(time);
-  const preview = document.createElement('div'); preview.className = 'hover-preview'; const previewCanvas = document.createElement('canvas'); previewCanvas.width = 256; previewCanvas.height = 256; if (state) { const context = previewCanvas.getContext('2d'); context.imageSmoothingEnabled = true; if (state instanceof HTMLCanvasElement) context.drawImage(state, 0, 0, 256, 256); else { const temporary = document.createElement('canvas'); temporary.width = SIZE; temporary.height = SIZE; drawState(temporary, state); context.drawImage(temporary, 0, 0, 256, 256); } } preview.append(previewCanvas);
-  canvas.addEventListener('mouseenter', () => { const rect = canvasWrap.getBoundingClientRect(); preview.style.left = `${Math.max(8, Math.min(window.innerWidth - 264, rect.left + rect.width / 2 - 128))}px`; preview.style.top = `${Math.max(8, Math.min(window.innerHeight - 264, rect.top + rect.height / 2 - 128))}px`; preview.classList.add('is-visible'); });
-  canvas.addEventListener('mouseleave', () => preview.classList.remove('is-visible'));
+  const preview = document.createElement('div'); preview.className = 'hover-preview';
+  const previewCanvas = document.createElement('canvas');
+  const nativeSize = modelInput.value === 'google/ddpm-cifar10-32' ? 32 : 256;
+  previewCanvas.width = state instanceof HTMLCanvasElement ? state.width : nativeSize;
+  previewCanvas.height = state instanceof HTMLCanvasElement ? state.height : nativeSize;
+  preview.style.width = `${previewCanvas.width + 4}px`; preview.style.height = `${previewCanvas.height + 4}px`;
+  previewCanvas.style.width = `${previewCanvas.width}px`; previewCanvas.style.height = `${previewCanvas.height}px`;
+  if (state) {
+    const context = previewCanvas.getContext('2d');
+    if (state instanceof HTMLCanvasElement) context.drawImage(state, 0, 0);
+    else { const temporary = document.createElement('canvas'); temporary.width = SIZE; temporary.height = SIZE; drawState(temporary, state); context.drawImage(temporary, 0, 0, previewCanvas.width, previewCanvas.height); }
+  }
+  preview.append(previewCanvas);
+  canvas.addEventListener('click', () => {
+    document.querySelectorAll('.hover-preview.is-visible').forEach(item => item.classList.remove('is-visible'));
+    const rect = canvasWrap.getBoundingClientRect();
+    preview.style.left = `${Math.max(8, Math.min(window.innerWidth - previewCanvas.width - 12, rect.left + rect.width / 2 - previewCanvas.width / 2))}px`;
+    preview.style.top = `${Math.max(8, Math.min(window.innerHeight - previewCanvas.height - 12, rect.top + rect.height / 2 - previewCanvas.height / 2))}px`;
+    preview.classList.add('is-visible');
+  });
   card.append(actionSlot, canvasWrap, caption, preview); return card;
 }
 function renderTracks() {
@@ -94,6 +111,7 @@ function renderTracks() {
   elements.denoisingTrack.style.setProperty('--stage-count', checkpoints.length);
   elements.diffusionTrack.replaceChildren(...checkpoints.map((step, index) => makeCard(step, forwardStates[index], 'forward', index)));
   elements.denoisingTrack.replaceChildren(...checkpoints.map((step, index) => makeCard(step, reverseStates.get(step) || (step === TOTAL_STEPS ? reverseNoiseState : null), 'reverse', index)));
+  modelInput.disabled = generationBusy;
   const button = document.querySelector('#denoise-button'); if (button) { button.disabled = generationBusy; button.textContent = generationBusy ? 'Denoising…' : 'Denoise'; }
   const forwardNoiseButton = document.querySelector('#forward-noise-button'); if (forwardNoiseButton) { forwardNoiseButton.disabled = generationBusy || !forwardCompleted; forwardNoiseButton.title = forwardCompleted ? 'Use the last forward-process image' : 'Run Diffuse first'; }
   const resampleButton = document.querySelector('#resample-button'); if (resampleButton) resampleButton.disabled = generationBusy;
@@ -144,7 +162,7 @@ async function generateReverse() {
   stopAnimation(); generationController = new AbortController(); generationBusy = true; reverseStates = new Map(); resetProgress('reverse'); setModelStatus('Loading model weights if needed, then sampling 1,000 reverse steps…'); renderTracks();
   try {
     const initialImage = reverseInputMode === 'forward' ? stateToPng(reverseNoiseState) : null;
-    const response = await fetch('/api/generate', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ checkpoints, seed, initial_image: initialImage }), signal: generationController.signal }); if (!response.ok || !response.body) throw new Error(`Backend request failed (${response.status}).`);
+    const response = await fetch('/api/generate', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ model: modelInput.value, checkpoints, seed, initial_image: initialImage }), signal: generationController.signal }); if (!response.ok || !response.body) throw new Error(`Backend request failed (${response.status}).`);
     const reader = response.body.getReader(); const decoder = new TextDecoder(); let pending = '';
     while (true) { const { value, done } = await reader.read(); pending += decoder.decode(value || new Uint8Array(), { stream: !done }); const lines = pending.split('\n'); pending = lines.pop(); for (const line of lines) if (line.trim()) await handleModelEvent(JSON.parse(line)); if (done) break; }
     if (pending.trim()) await handleModelEvent(JSON.parse(pending));
@@ -155,9 +173,20 @@ async function handleModelEvent(event) {
   if (event.type === 'status') setModelStatus(`Loading ${event.model} on the first generation…`);
   if (event.type === 'progress') { const completed = event.total - event.step; if (progressStartedAt.reverse === null) beginProgress('reverse'); setProgress('reverse', completed); setModelStatus(`Denoising · ${completed}/${event.total} steps complete.`); }
   if (event.type === 'frame') { reverseStates.set(event.step, await decodeFrame(event.image)); setProgress('reverse', TOTAL_STEPS - event.step); setModelStatus(`Real model inference · received t = ${event.step}.`); renderTracks(); const card = elements.denoisingTrack.querySelector(`[data-step="${event.step}"]`); if (card) { document.querySelectorAll('#denoising-track .is-active').forEach((item) => item.classList.remove('is-active')); card.classList.add('is-active'); } }
-  if (event.type === 'complete') { setProgress('reverse', TOTAL_STEPS); setModelStatus(`Generated with ${MODEL_NAME} on ${event.device.toUpperCase()} · seed ${event.seed}.`); } if (event.type === 'error') throw new Error(event.message);
+  if (event.type === 'complete') { setProgress('reverse', TOTAL_STEPS); setModelStatus(`Generated with ${event.model || modelInput.value} on ${event.device.toUpperCase()} · seed ${event.seed}.`); } if (event.type === 'error') throw new Error(event.message);
 }
 function loadImage(image, label) { const canvas = document.createElement('canvas'); canvas.width = SIZE; canvas.height = SIZE; const context = canvas.getContext('2d'); const scale = Math.max(SIZE / image.width, SIZE / image.height); const width = image.width * scale; const height = image.height * scale; context.drawImage(image, (SIZE - width) / 2, (SIZE - height) / 2, width, height); sourcePixels = canvasPixels(canvas); imageLabel = label; reverseInputMode = 'seed'; reverseNoiseState = makeNoiseState(readSeed()); forwardRevealCount = 1; forwardCompleted = false; resetProgress('forward'); computeForwardStates(); renderTracks(); }
+modelInput.addEventListener('change', () => {
+  applyTimesteps();
+  if (modelInput.value === 'google/ddpm-cifar10-32') {
+    const image = new Image();
+    image.onload = () => {
+      if (modelInput.value === 'google/ddpm-cifar10-32') loadImage(image, 'CIFAR-10 train 23');
+    };
+    image.onerror = () => setModelStatus('Could not load the CIFAR-10 source image.');
+    image.src = 'assets/cifar10_train_23_32.png';
+  }
+});
 elements.applyTimesteps.addEventListener('click', applyTimesteps); elements.timestepsInput.addEventListener('keydown', (event) => { if (event.key === 'Enter') applyTimesteps(); });
 elements.randomizeSeed.addEventListener('click', () => { resampleNoise(); elements.configError.textContent = ''; forwardRevealCount = 1; forwardCompleted = false; resetProgress('forward'); computeForwardStates(); renderTracks(); });
 elements.imageUpload.addEventListener('change', (event) => { const file = event.target.files?.[0]; if (!file) return; const image = new Image(); image.onload = () => { loadImage(image, file.name.replace(/\.[^.]+$/, '')); URL.revokeObjectURL(image.src); }; image.src = URL.createObjectURL(file); });
@@ -165,5 +194,14 @@ elements.imageUpload.addEventListener('change', (event) => { const file = event.
 // and as a fallback if the image asset cannot be read.
 sourcePixels = canvasPixels(makeSourceArt()); computeForwardStates(); reverseNoiseState = makeNoiseState(readSeed()); resetProgress('forward'); resetProgress('reverse'); renderTracks();
 const defaultSourceImage = new Image();
-defaultSourceImage.onload = () => loadImage(defaultSourceImage, 'Human face');
+defaultSourceImage.onload = () => {
+  if (modelInput.value === 'google/ddpm-celebahq-256') loadImage(defaultSourceImage, 'Human face');
+};
 defaultSourceImage.src = 'assets/celeba_hq_validation_91_256.png';
+
+document.addEventListener('pointerdown', event => {
+  if (!event.target.closest('.stage-canvas canvas')) document.querySelectorAll('.hover-preview.is-visible').forEach(item => item.classList.remove('is-visible'));
+});
+document.addEventListener('keydown', event => {
+  if (event.key === 'Escape') document.querySelectorAll('.hover-preview.is-visible').forEach(item => item.classList.remove('is-visible'));
+});

app.py

@@ -13,7 +13,7 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 
-from generate import CHECKPOINTS, ROOT, generate_images, load_model
+from generate import CHECKPOINTS, MODEL_ID, MODEL_SIZES, ROOT, generate_images, load_model
 
 ZERO_GPU = os.environ.get("SPACES_ZERO_GPU", "").lower() in {"1", "true", "yes"}
 def gpu_unavailable(error):
@@ -42,7 +42,8 @@ GPU_READY = False
 MPS_READY = False
 if ZERO_GPU:
     try:
-        load_model(device="cuda")
+        for model_id in MODEL_SIZES:
+            load_model(device="cuda", model_id=model_id)
         GPU_READY = True
     except Exception as error:
         if not gpu_unavailable(error):
@@ -76,8 +77,8 @@ def checkpoint_slots(checkpoints, frames=None):
     return [(frames.get(step, BLANK_IMAGE), f"t = {step}") for step in checkpoints]
 
 
-def gaussian_noise(seed):
-    noise = np.random.default_rng(int(seed)).standard_normal((256, 256, 3))
+def gaussian_noise(seed, model_id=MODEL_ID):
+    noise = np.random.default_rng(int(seed)).standard_normal((MODEL_SIZES[model_id], MODEL_SIZES[model_id], 3))
     return Image.fromarray(np.round((noise.clip(-1, 1) + 1) * 127.5).astype(np.uint8))
 
 
@@ -96,15 +97,15 @@ def progress_text(completed=0, started=None):
             f"{f'{rate:.2f}' if rate else '?'}it/s]")
 
 
-def diffuse(image, seed, checkpoint_text):
+def diffuse(image, seed, checkpoint_text, model_id=MODEL_ID):
     seed, checkpoints = parse_settings(seed, checkpoint_text)
     if image is None:
         raise gr.Error("Upload an image first.")
-    pixels = np.asarray(image.convert("RGB").resize((256, 256)), dtype=np.float32) / 127.5 - 1
+    pixels = np.asarray(image.convert("RGB").resize((MODEL_SIZES[model_id], MODEL_SIZES[model_id])), dtype=np.float32) / 127.5 - 1
     noise = np.random.default_rng(seed).standard_normal(pixels.shape)
     alpha_bars = np.cumprod(1 - np.linspace(0.0001, 0.02, 1000))
     started = time.monotonic()
-    frames = {0: image.convert("RGB").resize((256, 256))}
+    frames = {0: image.convert("RGB").resize((MODEL_SIZES[model_id], MODEL_SIZES[model_id]))}
     final_noise = None
     yield checkpoint_slots(checkpoints, frames), None, progress_html(0, started)
     for step in checkpoints:
@@ -116,12 +117,12 @@ def diffuse(image, seed, checkpoint_text):
         yield checkpoint_slots(checkpoints, frames), final_noise if step == 1000 else None, progress_html(step, started)
 
 
-def stream_denoising(seed, checkpoints, initial_image, device):
+def stream_denoising(seed, checkpoints, initial_image, device, model_id=MODEL_ID):
     started = time.monotonic()
-    frames = {1000: initial_image if initial_image is not None else gaussian_noise(seed)}
+    frames = {1000: initial_image if initial_image is not None else gaussian_noise(seed, model_id)}
     yield checkpoint_slots(checkpoints, frames), progress_html(0, started)
     print('using', device)
-    for step, frame in generate_images(seed, checkpoints, initial_image, device=device):
+    for step, frame in generate_images(seed, checkpoints, initial_image, device=device, model_id=model_id):
         if frame is not None:
             frames[step] = frame
         yield (checkpoint_slots(checkpoints, frames) if frame is not None else gr.skip(),
@@ -129,18 +130,18 @@ def stream_denoising(seed, checkpoints, initial_image, device):
 
 
 @spaces.GPU(duration=int(os.environ.get("GPU_DURATION_SECONDS", "120")))
-def denoise_gpu(seed, checkpoints, initial_image):
-    yield from stream_denoising(seed, checkpoints, initial_image, "cuda")
+def denoise_gpu(seed, checkpoints, initial_image, model_id=MODEL_ID):
+    yield from stream_denoising(seed, checkpoints, initial_image, "cuda", model_id)
 
 
-def denoise(seed, checkpoint_text, input_mode, forward_noise):
+def denoise(seed, checkpoint_text, input_mode, forward_noise, model_id=MODEL_ID):
     seed, checkpoints = parse_settings(seed, checkpoint_text)
     initial_image = forward_noise if input_mode == "Forward noise" else None
     if input_mode == "Forward noise" and initial_image is None:
         raise gr.Error("Run Diffuse first to create forward noise.")
     if GPU_READY:
         try:
-            yield from denoise_gpu(seed, checkpoints, initial_image)
+            yield from denoise_gpu(seed, checkpoints, initial_image, model_id)
             return
         except Exception as error:
             if not gpu_unavailable(error):
@@ -150,7 +151,7 @@ def denoise(seed, checkpoint_text, input_mode, forward_noise):
             gr.Info(f"CUDA unavailable or quota exhausted. Running on {fallback}.")
     if MPS_READY:
         try:
-            yield from stream_denoising(seed, checkpoints, initial_image, "mps")
+            yield from stream_denoising(seed, checkpoints, initial_image, "mps", model_id)
             return
         except (RuntimeError, NotImplementedError) as error:
             message = str(error).lower()
@@ -158,7 +159,7 @@ def denoise(seed, checkpoint_text, input_mode, forward_noise):
                 raise
             print(f"MPS unavailable; restarting on CPU: {error}")
             gr.Info("MPS unavailable. Running on CPU.")
-    yield from stream_denoising(seed, checkpoints, initial_image, "cpu")
+    yield from stream_denoising(seed, checkpoints, initial_image, "cpu", model_id)
 
 
 # Reuse the original page's palette, typography, header, and panel styling.
@@ -201,7 +202,23 @@ PAGE_CSS = (ROOT / "styles.css").read_text() + """
 .config-field textarea,.config-field input { width:100% !important; max-width:none !important; box-sizing:border-box !important; height:40px !important; min-height:40px !important;
   border:1px solid #cfd3c9 !important; border-radius:0 3px 3px 0 !important;
   padding:10px 12px !important; font:11px var(--mono) !important; background:white !important; white-space:pre !important; overflow-x:auto !important; resize:none !important; }
-#model-field textarea { background:#f8f7f2 !important; color:#687169 !important; }
+/* The dropdown wrapper adds spacing that textboxes do not have. */
+.config-row:has(#model-field) { column-gap:0 !important; }
+.config-row:has(#model-field) > .form { margin:0 !important; padding:0 !important; gap:0 !important; }
+#model-field { margin:0 !important; padding:0 !important;
+  height:40px !important; min-height:40px !important; max-height:40px !important; }
+#model-field .wrap, #model-field .input-container { margin:0 !important; }
+#model-field .container { padding:0 !important; margin:0 !important; }
+#model-field .wrap { box-sizing:border-box !important;
+  height:40px !important; min-height:40px !important; max-height:40px !important;
+  border:1px solid #cfd3c9 !important;
+  border-radius:0 3px 3px 0 !important; box-shadow:none !important; }
+#model-field .wrap-inner { height:38px !important; min-height:0 !important;
+  padding:0 12px !important; gap:0 !important; }
+/* The dropdown input sits inside its own bordered wrapper. */
+#model-field input { border:0 !important; border-radius:0 !important;
+  margin:0 !important; padding:0 !important; background:transparent !important; box-shadow:none !important;
+  height:38px !important; min-height:38px !important; }
 .process-button { flex:0 0 auto !important; min-width:76px !important;
   height:28px !important; min-height:28px !important; padding:0 12px !important;
   background:var(--green) !important; color:white !important; border-radius:3px !important;
@@ -234,10 +251,10 @@ PAGE_CSS = (ROOT / "styles.css").read_text() + """
 .checkpoint-preview-host { position:relative; height:0 !important; min-height:0 !important; padding:0 !important;
   margin:0 !important; border:0 !important; overflow:visible !important; z-index:1000; }
 .checkpoint-preview-host .html-container { padding:0 !important; overflow:visible !important; }
-.checkpoint-preview { position:absolute; bottom:8px; width:min(256px,100%); z-index:1000;
+.checkpoint-preview { position:absolute; bottom:8px; width:max-content; max-width:none; z-index:1000;
   border:2px solid #fbfaf6; background:#fbfaf6; box-shadow:0 8px 24px #25312852; }
 .checkpoint-preview summary { cursor:pointer; padding:3px 6px; font:10px var(--mono); }
-.checkpoint-preview img { width:100%; height:auto; display:block; }
+.checkpoint-preview img { width:auto; height:auto; max-width:none; display:block; }
 .checkpoint-preview:not([open]) { display:none; }
 #reverse-input { justify-self:end; width:max-content !important; min-width:0 !important;
   border:0 !important; background:transparent !important; padding:0 !important; }
@@ -287,19 +304,26 @@ def progress_html(completed=0, started=None):
     return "<pre>" + progress_text(completed, started) + "</pre>"
 
 
-def reset_tracks(seed, checkpoint_text, image):
+def reset_tracks(seed, checkpoint_text, image, model_id=MODEL_ID):
     seed, steps = parse_settings(seed, checkpoint_text)
-    initial = {0: image.convert("RGB").resize((256, 256))} if image is not None else {}
+    initial = {0: image.convert("RGB").resize((MODEL_SIZES[model_id], MODEL_SIZES[model_id]))} if image is not None else {}
     return (", ".join(map(str, steps)),
             gr.Gallery(value=checkpoint_slots(steps, initial), columns=len(steps), allow_preview=False, interactive=False),
-            gr.Gallery(value=checkpoint_slots(steps, {1000: gaussian_noise(seed)}), columns=len(steps), allow_preview=False, interactive=False),
+            gr.Gallery(value=checkpoint_slots(steps, {1000: gaussian_noise(seed, model_id)}), columns=len(steps), allow_preview=False, interactive=False),
             None, progress_html(), progress_html())
 
 
-def change_source(filepath, seed, checkpoint_text):
+def change_source(filepath, seed, checkpoint_text, model_id=MODEL_ID):
     with Image.open(filepath) as uploaded:
         image = uploaded.convert("RGB")
-    return (image, *reset_tracks(seed, checkpoint_text, image))
+    return (image, *reset_tracks(seed, checkpoint_text, image, model_id))
+
+
+def change_model(seed, checkpoint_text, image, model_id):
+    if model_id == "google/ddpm-cifar10-32":
+        return change_source(ROOT / "assets/cifar10_train_23_32.png",
+                             seed, checkpoint_text, model_id)
+    return (image, *reset_tracks(seed, checkpoint_text, image, model_id))
 
 
 def show_checkpoint_preview(gallery, event: gr.SelectData):
@@ -315,14 +339,16 @@ def show_checkpoint_preview(gallery, event: gr.SelectData):
     buffer = BytesIO()
     image.save(buffer, format="PNG")
     encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+    width, height = image.size
+    preview_width = width + 4
     center = 100 * (index + 0.5) / len(gallery)
     # Force fresh markup even when the same thumbnail is clicked again, so
     # a details element closed in the browser is recreated in its open state.
     preview_id = secrets.token_hex(8)
     return (f'<details open class="checkpoint-preview" data-preview-id="{preview_id}" ' 
-            f'style="left:clamp(0px, calc({center}% - 128px), max(0px, calc(100% - 256px)))">'
+            f'style="left:clamp(0px, calc({center}% - {preview_width / 2}px), max(0px, calc(100% - {preview_width}px)))">'
             '<summary>Close preview</summary>'
-            f'<img src="data:image/png;base64,{encoded}" alt="{html.escape(caption or "Checkpoint", quote=True)}">'
+            f'<img width="{width}" height="{height}" src="data:image/png;base64,{encoded}" alt="{html.escape(caption or "Checkpoint", quote=True)}">'
             '</details>')
 
 
@@ -337,8 +363,8 @@ with gr.Blocks(title="Diffusion — A DDPM in motion", **BLOCK_OPTIONS) as demo:
             gr.HTML("<h1>Configure</h1>")
             with gr.Row(elem_classes="config-row"):
                 gr.Button("Model", interactive=False, scale=0, min_width=110, elem_classes="config-label")
-                gr.Textbox(value="google/ddpm-celebahq-256", interactive=False,
-                           scale=10, min_width=0, lines=1, max_lines=1, show_label=False, container=False, elem_id="model-field", elem_classes="config-field")
+                model = gr.Dropdown(choices=list(MODEL_SIZES), value=MODEL_ID,
+                           scale=10, min_width=0, show_label=False, container=False, elem_id="model-field", elem_classes="config-field")
             with gr.Row(elem_classes="config-row"):
                 gr.Button("Timesteps", interactive=False, scale=0, min_width=110, elem_classes="config-label")
                 checkpoints = gr.Textbox(value=", ".join(map(str, CHECKPOINTS)), show_label=False,
@@ -373,11 +399,11 @@ with gr.Blocks(title="Diffusion — A DDPM in motion", **BLOCK_OPTIONS) as demo:
             forward_gallery.select(show_checkpoint_preview, [forward_gallery], [forward_preview],
                                   queue=False, show_progress="hidden").then(
                                       fn=None, js=INSTALL_PREVIEW_DISMISS, queue=False)
-            diffuse_button.click(diffuse, [source, seed, checkpoints],
+            diffuse_button.click(diffuse, [source, seed, checkpoints, model],
                                  [forward_gallery, forward_noise, forward_status])
 
         with gr.Column(elem_classes=["lab-panel", "process-panel"]):
-            gr.HTML('<div class="process-heading"><h2>Reverse process: noise to face</h2>'
+            gr.HTML('<div class="process-heading"><h2>Reverse process: noise to image</h2>'
                     '<p class="reverse-equation">x<sub>t−1</sub> = 1/√α<sub>t</sub>'
                     ' (x<sub>t</sub> − β<sub>t</sub>/√(1−α̅<sub>t</sub>)'
                     ' ε<sub>θ</sub>(x<sub>t</sub>,t)) + σ<sub>t</sub>z</p></div>')
@@ -395,7 +421,7 @@ with gr.Blocks(title="Diffusion — A DDPM in motion", **BLOCK_OPTIONS) as demo:
             reverse_gallery.select(show_checkpoint_preview, [reverse_gallery], [reverse_preview],
                                   queue=False, show_progress="hidden").then(
                                       fn=None, js=INSTALL_PREVIEW_DISMISS, queue=False)
-            denoise_button.click(denoise, [seed, checkpoints, input_mode, forward_noise],
+            denoise_button.click(denoise, [seed, checkpoints, input_mode, forward_noise, model],
                                  [reverse_gallery, status], concurrency_limit=1,
                                  **({"stream_every": 0} if "stream_every" in
                                     inspect.signature(denoise_button.click).parameters else {}))
@@ -408,8 +434,10 @@ with gr.Blocks(title="Diffusion — A DDPM in motion", **BLOCK_OPTIONS) as demo:
                                          queue=False, show_progress="hidden")
         reset_outputs = [checkpoints, forward_gallery, reverse_gallery,
                          forward_noise, forward_status, status]
-        apply_button.click(reset_tracks, [seed, checkpoints, source], reset_outputs, queue=False)
-        upload_source.upload(change_source, [upload_source, seed, checkpoints],
+        apply_button.click(reset_tracks, [seed, checkpoints, source, model], reset_outputs, queue=False)
+        model.change(change_model, [seed, checkpoints, source, model], [source, *reset_outputs])
+        model.change(clear_previews, outputs=[forward_preview, reverse_preview], queue=False)
+        upload_source.upload(change_source, [upload_source, seed, checkpoints, model],
                              [source, *reset_outputs], queue=False)
 
 if __name__ == "__main__":

download_cifar10.py

@@ -0,0 +1,83 @@
+import argparse
+import json
+from io import BytesIO
+from pathlib import Path
+from urllib.error import HTTPError, URLError
+from urllib.parse import urlencode
+from urllib.request import Request, urlopen
+
+from PIL import Image
+
+
+DATASET = "uoft-cs/cifar10"
+API_URL = "https://datasets-server.huggingface.co/rows"
+
+
+def download_row(row_index, split="test", output=None):
+    """Download one CIFAR-10 row and save it as a 32x32 RGB PNG."""
+    if row_index < 0:
+        raise ValueError("Row index must be zero or greater.")
+
+    if split not in {"train", "test"}:
+        raise ValueError("Split must be train or test.")
+
+    query = urlencode(
+        {
+            "dataset": DATASET,
+            "config": "plain_text",
+            "split": split,
+            "offset": row_index,
+            "length": 1,
+        }
+    )
+    request = Request(f"{API_URL}?{query}", headers={"User-Agent": "diffusion-lab/1.0"})
+    with urlopen(request, timeout=30) as response:
+        result = json.load(response)
+
+    rows = result.get("rows", [])
+    if not rows:
+        total = result.get("num_rows_total", "the available")
+        raise ValueError(f"Row {row_index} is outside {split!r} ({total} rows).")
+
+    image_url = rows[0]["row"]["img"]["src"]
+    image_request = Request(image_url, headers={"User-Agent": "diffusion-lab/1.0"})
+    with urlopen(image_request, timeout=60) as response:
+        image_bytes = response.read()
+
+    destination = Path(output or f"cifar10_{split}_{row_index}_32.png")
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with Image.open(BytesIO(image_bytes)) as image:
+        image = image.convert("RGB").resize((32, 32), Image.Resampling.LANCZOS)
+        image.save(destination, format="PNG", optimize=True)
+
+    return destination
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Download a specific row from the CIFAR-10 32x32 dataset."
+    )
+    parser.add_argument("row_index", type=int, help="Zero-based row index")
+    parser.add_argument(
+        "--split",
+        choices=("train", "test"),
+        default="test",
+        help="Dataset split (default: test)",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        help="Output PNG path (default: cifar10_<split>_<row>_32.png)",
+    )
+    args = parser.parse_args()
+
+    try:
+        destination = download_row(args.row_index, args.split, args.output)
+    except (HTTPError, URLError, TimeoutError, ValueError, KeyError, OSError) as error:
+        parser.exit(1, f"Download failed: {error}\n")
+
+    print(f"Saved row {args.row_index} from {args.split!r} to {destination}")
+
+
+if __name__ == "__main__":
+    main()

generate.py

@@ -9,6 +9,7 @@ from PIL import Image
 
 ROOT = Path(__file__).resolve().parent
 MODEL_ID = "google/ddpm-celebahq-256"
+MODEL_SIZES = {MODEL_ID: 256, "google/ddpm-cifar10-32": 32}
 CHECKPOINTS = (0, 50, 100, 200, 350, 500, 700, 1000)
 
 # Gradio may resume a generator on a different worker thread between yields.
@@ -21,10 +22,12 @@ _TORCH = None
 _DEVICE = None
 
 
-def load_model(device=None):
+def load_model(device=None, model_id=MODEL_ID):
     """Load and cache the diffusion pipeline, returning it and its device."""
     global _PIPELINE, _TORCH, _DEVICE
 
+    if model_id not in MODEL_SIZES:
+        raise ValueError(f"Unsupported model: {model_id}")
     with _CACHE_LOCK:
         try:
             import torch
@@ -44,13 +47,14 @@ def load_model(device=None):
             else:
                 device = "cpu"
         device = str(device)
-        if device not in _PIPELINES:
-            print(f"Loading model on device: {device}")
+        key = (model_id, device)
+        if key not in _PIPELINES:
+            print(f"Loading {model_id} on device: {device}")
             pipeline = DDPMPipeline.from_pretrained(
-                MODEL_ID, cache_dir=Path(os.environ.get("MODEL_CACHE_DIR", ROOT / "models")))
+                model_id, cache_dir=Path(os.environ.get("MODEL_CACHE_DIR", ROOT / "models")))
             pipeline.to(device)
-            _PIPELINES[device] = pipeline
-        _PIPELINE = _PIPELINES[device]
+            _PIPELINES[key] = pipeline
+        _PIPELINE = _PIPELINES[key]
         _TORCH = torch
         _DEVICE = device
         return _PIPELINE, _TORCH, _DEVICE
@@ -67,10 +71,10 @@ def sample_to_image(sample):
     return Image.fromarray(image, mode="RGB")
 
 
-def generate_images(seed, checkpoints=CHECKPOINTS, initial_image=None, device=None):
+def generate_images(seed, checkpoints=CHECKPOINTS, initial_image=None, device=None, model_id=MODEL_ID):
     """Yield every remaining-step count and an image at requested checkpoints."""
     with _MODEL_LOCK:
-        pipeline, torch, device = load_model(device=device)
+        pipeline, torch, device = load_model(device=device, model_id=model_id)
         scheduler = pipeline.scheduler
         training_steps = scheduler.config.num_train_timesteps
         scheduler.set_timesteps(training_steps, device=device)
@@ -117,10 +121,10 @@ def generate_images(seed, checkpoints=CHECKPOINTS, initial_image=None, device=No
                 yield remaining_steps, image
 
 
-def generate_image(seed):
+def generate_image(seed, model_id=MODEL_ID):
     """Generate and return the final image for a seed."""
     final_image = None
-    for _, image in generate_images(seed, checkpoints=(0,)):
+    for _, image in generate_images(seed, checkpoints=(0,), model_id=model_id):
         if image is not None:
             final_image = image
     if final_image is None:
@@ -129,25 +133,26 @@ def generate_image(seed):
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Generate a CelebA-HQ face image.")
+    parser = argparse.ArgumentParser(description="Generate an unconditional DDPM image.")
     parser.add_argument(
         "--seed",
         type=int,
         default=None,
         help="Random seed (defaults to a randomly generated seed)",
     )
+    parser.add_argument("--model", choices=list(MODEL_SIZES), default=MODEL_ID)
     args = parser.parse_args()
     seed = args.seed if args.seed is not None else secrets.randbelow(2**31)
     if seed < 0 or seed >= 2**63:
         parser.error("seed must be between 0 and 2^63 - 1")
 
     print(f"Generating with seed: {seed}")
-    image = generate_image(seed)
+    image = generate_image(seed, model_id=args.model)
 
     import matplotlib.pyplot as plt
 
     plt.imshow(image)
-    plt.title(f"{MODEL_ID} — seed {seed}")
+    plt.title(f"{args.model} — seed {seed}")
     plt.axis("off")
     plt.tight_layout()
     plt.show()

index.html

@@ -23,7 +23,10 @@
         <div class="config-rows">
           <div class="config-row">
             <label for="model-input">Model</label>
-            <input id="model-input" type="text" value="google/ddpm-celebahq-256" readonly>
+            <select id="model-input">
+              <option value="google/ddpm-celebahq-256">google/ddpm-celebahq-256</option>
+              <option value="google/ddpm-cifar10-32">google/ddpm-cifar10-32</option>
+            </select>
           </div>
           <div class="config-row">
             <label for="timesteps-input">Timesteps</label>
@@ -48,7 +51,7 @@
       </section>
       <section class="process-panel lab-panel" aria-labelledby="denoising-title">
         <div class="process-heading"><div class="panel-heading"><div>
-							<h2 id="denoising-title">Reverse process: noise to face</h2></div></div>
+							<h2 id="denoising-title">Reverse process: noise to image</h2></div></div>
           <p class="reverse-equation">x<sub>t−1</sub> = 1/√α<sub>t</sub> (x<sub>t</sub> − β<sub>t</sub>/√(1−α̅<sub>t</sub>) ε<sub>θ</sub>(x<sub>t</sub>,t)) + σ<sub>t</sub>z</p>
         </div>
         <output id="reverse-progress-label" class="step-progress">0/1000</output>

server.py

@@ -8,7 +8,7 @@ from io import BytesIO
 from pathlib import Path
 from urllib.parse import urlparse
 
-from generate import CHECKPOINTS, MODEL_ID, generate_images, is_model_loaded, load_model
+from generate import CHECKPOINTS, MODEL_ID, MODEL_SIZES, generate_images, is_model_loaded, load_model
 
 
 ROOT = Path(__file__).resolve().parent
@@ -20,17 +20,17 @@ def encode_frame(image):
     return base64.b64encode(buffer.getvalue()).decode("ascii")
 
 
-def generate_frames(seed, checkpoints=CHECKPOINTS, initial_image=None):
-    yield {"type": "status", "status": "loading", "model": MODEL_ID}
+def generate_frames(seed, checkpoints=CHECKPOINTS, initial_image=None, model_id=MODEL_ID):
+    yield {"type": "status", "status": "loading", "model": model_id}
 
     try:
-        for step, image in generate_images(seed, checkpoints=checkpoints, initial_image=initial_image):
+        for step, image in generate_images(seed, checkpoints=checkpoints, initial_image=initial_image, model_id=model_id):
             yield {"type": "progress", "step": step, "total": 1000}
             if image is not None:
                 yield {"type": "frame", "step": step, "image": encode_frame(image)}
 
-        _, _, device = load_model()
-        yield {"type": "complete", "seed": seed, "device": device}
+        _, _, device = load_model(model_id=model_id)
+        yield {"type": "complete", "seed": seed, "device": device, "model": model_id}
     except Exception as error:
         yield {"type": "error", "message": str(error)}
 
@@ -47,7 +47,7 @@ class DiffusionHandler(BaseHTTPRequestHandler):
     def do_GET(self):
         path = urlparse(self.path).path
         if path == "/api/health":
-            body = json.dumps({"model": MODEL_ID, "loaded": is_model_loaded()}).encode()
+            body = json.dumps({"model": MODEL_ID, "models": list(MODEL_SIZES), "loaded": is_model_loaded()}).encode()
             self.send_bytes(200, "application/json", body)
             return
 
@@ -57,6 +57,7 @@ class DiffusionHandler(BaseHTTPRequestHandler):
             "app.js",
             "styles.css",
             "assets/celeba_hq_validation_91_256.png",
+            "assets/cifar10_train_23_32.png",
         }:
             self.send_error(404)
             return
@@ -73,6 +74,9 @@ class DiffusionHandler(BaseHTTPRequestHandler):
         try:
             length = int(self.headers.get("Content-Length", "0"))
             request = json.loads(self.rfile.read(length) or b"{}")
+            model_id = request.get("model", MODEL_ID)
+            if not isinstance(model_id, str) or model_id not in MODEL_SIZES:
+                raise ValueError("Unsupported model.")
             seed = int(request.get("seed", secrets.randbelow(2**31)))
             if seed < 0 or seed >= 2**63:
                 raise ValueError("Seed must be between 0 and 2^63 - 1.")
@@ -106,7 +110,7 @@ class DiffusionHandler(BaseHTTPRequestHandler):
         self.send_header("X-Accel-Buffering", "no")
         self.end_headers()
         try:
-            for event in generate_frames(seed, checkpoints, initial_image):
+            for event in generate_frames(seed, checkpoints, initial_image, model_id):
                 self.wfile.write((json.dumps(event) + "\n").encode())
                 self.wfile.flush()
         except (BrokenPipeError, ConnectionResetError):

styles.css

@@ -23,8 +23,8 @@ h1,h2 { margin:0; font:600 20px var(--display); } h1{font-size:18px}
 .config-rows { max-width:820px; display:grid; gap:8px; }
 .config-row { display:grid; grid-template-columns:110px minmax(0,1fr) 70px; align-items:center; min-height:40px; }
 .config-row label,.row-label-button { align-self:stretch; display:flex; align-items:center; padding:0 12px; border:1px solid #cfd3c9; border-right:0; border-radius:3px 0 0 3px; color:var(--green); background:#f0f2eb; font:10px var(--mono); }
-.config-row input { min-width:0; width:100%; height:40px; padding:0 12px; color:var(--ink); border:1px solid #cfd3c9; border-radius:0 3px 3px 0; outline:none; background:white; font:11px var(--mono); }
-.config-row input:focus { position:relative; border-color:var(--green); box-shadow:inset 0 0 0 1px var(--green); }
+.config-row input, .config-row select { min-width:0; width:100%; height:40px; padding:0 12px; color:var(--ink); border:1px solid #cfd3c9; border-radius:0 3px 3px 0; outline:none; background:white; font:11px var(--mono); }
+.config-row input:focus, .config-row select:focus { position:relative; border-color:var(--green); box-shadow:inset 0 0 0 1px var(--green); }
 .config-row input[readonly] { color:#687169; background:#f8f7f2; }
 .config-row #model-input,.config-row #seed-input { grid-column:2 / 4; }
 .config-row:has(#apply-timesteps) input { border-radius:0; }