| from __future__ import annotations |
|
|
| import base64 |
| import io |
| import os |
| import time |
| from functools import lru_cache |
|
|
| import gradio as gr |
| import numpy as np |
| import onnxruntime as ort |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
|
|
| try: |
| import spaces |
| except ImportError: |
| class _LocalSpaces: |
| """Keep the Space runnable outside Hugging Face ZeroGPU.""" |
|
|
| @staticmethod |
| def GPU(*args, **kwargs): |
| del args, kwargs |
|
|
| def decorator(function): |
| return function |
|
|
| return decorator |
|
|
| spaces = _LocalSpaces() |
|
|
|
|
| MODEL_REPO = os.getenv("MODEL_REPO", "Elixpo/monodiffuse") |
| MODEL_DIR = os.getenv("MODEL_DIR") |
| DIFFUSION_STEPS = 1000 |
| FRAME_DELAY_SECONDS = 0.16 |
| TARGET_PREVIEW_FRAMES = 30 |
| FP32_DIGIT_SEEDS = {0: 15, 1: 16, 2: 50, 3: 22, 4: 0, 5: 0, 6: 12, 7: 4, 8: 10, 9: 18} |
| MODEL_FILES = { |
| "FP32 deployment": "mnist_conditional_fp32.onnx", |
| "Native W1A32 (experimental)": "mnist_conditional_w1a32.onnx", |
| } |
|
|
|
|
| def session_options() -> ort.SessionOptions: |
| options = ort.SessionOptions() |
| options.intra_op_num_threads = max(1, min(4, os.cpu_count() or 1)) |
| options.inter_op_num_threads = 1 |
| options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| return options |
|
|
|
|
| @lru_cache(maxsize=3) |
| def load_session(filename: str) -> ort.InferenceSession: |
| path = os.path.join(MODEL_DIR, filename) if MODEL_DIR else hf_hub_download( |
| repo_id=MODEL_REPO, filename=filename |
| ) |
| return ort.InferenceSession(path, session_options(), providers=["CPUExecutionProvider"]) |
|
|
|
|
| def display_frame(state: np.ndarray) -> str: |
| pixels = np.clip((state.squeeze() + 1.0) * 127.5, 0, 255).astype(np.uint8) |
| image = Image.fromarray(pixels, mode="L") |
| buffer = io.BytesIO() |
| image.save(buffer, format="PNG", optimize=True) |
| encoded = base64.b64encode(buffer.getvalue()).decode("ascii") |
| return ( |
| '<div class="preview-frame">' |
| f'<img src="data:image/png;base64,{encoded}" alt="Current diffusion frame">' |
| "</div>" |
| ) |
|
|
|
|
| def classifier_prediction(state: np.ndarray) -> tuple[int, float]: |
| session = load_session("mnist_classifier.onnx") |
| logits = session.run(None, {"image": np.asarray(state, dtype=np.float32)})[0][0] |
| probabilities = np.exp(logits - logits.max()) |
| probabilities /= probabilities.sum() |
| digit = int(probabilities.argmax()) |
| return digit, float(probabilities[digit]) |
|
|
|
|
| def sampling_plan(mode: str) -> tuple[str, np.ndarray]: |
| if mode == "Full DDPM · 1000 steps": |
| return "ddpm", np.arange(DIFFUSION_STEPS - 1, -1, -1, dtype=np.int64) |
| steps = int(mode.split("·", 1)[1].split()[0]) |
| return "ddim", np.linspace(DIFFUSION_STEPS - 1, 0, steps, dtype=np.int64) |
|
|
|
|
| @spaces.GPU(duration=30) |
| def generate(digit: int, model_name: str, mode: str, seed: int, curated_seed: bool): |
| digit = int(digit) |
| seed = FP32_DIGIT_SEEDS[digit] if curated_seed and model_name == "FP32 deployment" else int(seed) |
| sampler, schedule = sampling_plan(mode) |
| session = load_session(MODEL_FILES[model_name]) |
|
|
| betas = np.linspace(1e-4, 0.02, DIFFUSION_STEPS, dtype=np.float32) |
| alphas = 1.0 - betas |
| alpha_bars = np.cumprod(alphas, dtype=np.float32) |
| previous = np.concatenate((np.ones(1, dtype=np.float32), alpha_bars[:-1])) |
| posterior_variance = betas * (1.0 - previous) / (1.0 - alpha_bars) |
|
|
| rng = np.random.default_rng(seed) |
| state = rng.standard_normal((1, 1, 28, 28), dtype=np.float32) |
| emit_every = max(1, (len(schedule) + TARGET_PREVIEW_FRAMES - 1) // TARGET_PREVIEW_FRAMES) |
| started = time.perf_counter() |
|
|
| yield display_frame(np.clip(state, -1, 1)), ( |
| f"### Noise initialized\nDigit **{digit}** · seed **{seed}** · {model_name}" |
| ) |
| time.sleep(FRAME_DELAY_SECONDS) |
|
|
| for position, timestep_value in enumerate(schedule): |
| timestep = int(timestep_value) |
| inputs = { |
| "image": state, |
| "timestep": np.asarray([timestep], dtype=np.int64), |
| "label": np.asarray([digit], dtype=np.int64), |
| } |
| predicted_noise = session.run(None, inputs)[0] |
| alpha_bar = alpha_bars[timestep] |
| predicted_clean = ( |
| state - np.sqrt(1.0 - alpha_bar) * predicted_noise |
| ) / np.sqrt(alpha_bar) |
|
|
| if sampler == "ddpm": |
| mean = state / np.sqrt(alphas[timestep]) - ( |
| betas[timestep] |
| / (np.sqrt(alphas[timestep]) * np.sqrt(1.0 - alpha_bar)) |
| ) * predicted_noise |
| if timestep: |
| noise = rng.standard_normal(state.shape, dtype=np.float32) |
| state = mean + np.sqrt(posterior_variance[timestep]) * noise |
| else: |
| state = mean |
| elif position == len(schedule) - 1: |
| state = predicted_clean |
| else: |
| next_alpha_bar = alpha_bars[int(schedule[position + 1])] |
| state = ( |
| np.sqrt(next_alpha_bar) * predicted_clean |
| + np.sqrt(1.0 - next_alpha_bar) * predicted_noise |
| ).astype(np.float32, copy=False) |
|
|
| completed = position + 1 |
| if completed % emit_every == 0 or completed == len(schedule): |
| elapsed = time.perf_counter() - started |
| percent = 100.0 * completed / len(schedule) |
| yield display_frame(np.clip(predicted_clean, -1, 1)), ( |
| f"### Denoising · {percent:.0f}%\n" |
| f"Step **{completed}/{len(schedule)}** · diffusion timestep **{timestep}** · " |
| f"elapsed **{elapsed:.1f}s**" |
| ) |
| time.sleep(FRAME_DELAY_SECONDS) |
|
|
| state = np.clip(state, -1, 1).astype(np.float32, copy=False) |
| predicted_digit, confidence = classifier_prediction(state) |
| elapsed = time.perf_counter() - started |
| match = predicted_digit == digit |
| badge = "✅ match" if match else "⚠️ mismatch" |
| yield display_frame(state), ( |
| f"### Complete · {badge}\nRequested **{digit}** · evaluator prediction " |
| f"**{predicted_digit}** · confidence **{confidence:.1%}** · seed **{seed}** · " |
| f"**{elapsed:.1f}s**" |
| ) |
|
|
|
|
| CSS = """ |
| .gradio-container {max-width: 1080px !important; margin: auto !important; |
| padding: 10px 18px 8px !important; min-height: 100vh;} |
| #hero {text-align:center; padding: 0 0 6px;} |
| #hero h1 {font-size: clamp(2rem, 4vw, 3.1rem); line-height: 1; margin: 0 0 .12em; |
| background: linear-gradient(90deg,#7c3aed,#22d3ee); -webkit-background-clip:text; |
| -webkit-text-fill-color:transparent;} |
| #hero p {margin: 0; font-size: .95rem; opacity: .78;} |
| #workspace {gap: 14px !important; align-items: stretch;} |
| #stage {border: 1px solid rgba(124,58,237,.35); border-radius: 24px; |
| box-shadow: 0 20px 80px rgba(76,29,149,.16); overflow:hidden;} |
| #preview {height: 330px; min-height: 330px; padding: 0 !important; |
| background: radial-gradient(circle at 50% 45%,rgba(124,58,237,.12),transparent 58%);} |
| #preview .preview-frame {height: 100%; width: 100%; display: flex; |
| align-items: center; justify-content: center;} |
| #preview .preview-frame img {width: min(92%,330px); height: min(92%,330px); |
| object-fit: contain; image-rendering: pixelated; image-rendering: crisp-edges;} |
| #status {min-height: 64px; padding: 2px 12px; font-size: .9rem;} |
| #status h3 {margin: .2rem 0 !important; font-size: 1rem;} |
| #controls {gap: 8px !important;} |
| #controls .block {margin: 0 !important;} |
| .compact-row {gap: 8px !important;} |
| .gradient-selector .wrap {gap: 6px !important; flex-direction: row !important; |
| flex-wrap: nowrap !important;} |
| .gradient-selector label:has(input[type="radio"]) {flex: 1 1 0; min-width: 0; |
| justify-content: center; text-align: center; padding: 8px 7px !important; |
| border: 1px solid rgba(124,58,237,.3) !important; border-radius: 11px !important; |
| background: linear-gradient(145deg,rgba(124,58,237,.08),rgba(34,211,238,.05)); |
| transition: border-color .15s ease, transform .15s ease, box-shadow .15s ease;} |
| .gradient-selector label:has(input[type="radio"]):hover {transform: translateY(-1px); |
| border-color: rgba(34,211,238,.8) !important;} |
| .gradient-selector label:has(input[type="radio"]:checked) { |
| color: white !important; border-color: transparent !important; |
| background: linear-gradient(115deg,#7c3aed,#2563eb 55%,#0891b2) !important; |
| box-shadow: 0 5px 16px rgba(76,29,149,.25);} |
| .gradient-selector label:has(input[type="radio"]:checked) span {color: white !important;} |
| .gradient-selector input[type="radio"] {display: none !important;} |
| #run-button {min-height: 42px !important; border: 0 !important; |
| background: linear-gradient(105deg,#7c3aed,#2563eb 55%,#0891b2) !important; |
| box-shadow: 0 8px 22px rgba(76,29,149,.24);} |
| #footer-note {font-size: .78rem; line-height: 1.35; opacity: .72; text-align: center; |
| padding: 3px 10px 0;} |
| #footer-note p {margin: 0 !important;} |
| @media (max-width: 760px) { |
| .gradio-container {padding: 8px 10px !important;} |
| .gradient-selector .wrap {flex-wrap: wrap !important;} |
| .gradient-selector label:has(input[type="radio"]) {min-width: 46%;} |
| } |
| """ |
|
|
|
|
| with gr.Blocks() as demo: |
| gr.HTML( |
| "<div id='hero'><h1>Diffuino Live</h1></div>" |
| ) |
| with gr.Row(equal_height=True, elem_id="workspace"): |
| with gr.Column(scale=2, elem_id="stage"): |
| output = gr.HTML( |
| '<div class="preview-frame"><span>Ready for diffusion</span></div>', |
| elem_id="preview", |
| ) |
| status = gr.Markdown("### Ready\nChoose a digit and start diffusion.", elem_id="status") |
| with gr.Column(scale=1, elem_id="controls"): |
| with gr.Row(elem_classes=["compact-row"]): |
| digit = gr.Slider(0, 9, value=6, step=1, label="Digit", scale=3) |
| seed = gr.Number(value=12, precision=0, label="Seed", scale=1) |
| model = gr.Radio( |
| list(MODEL_FILES), |
| value="FP32 deployment", |
| label="Checkpoint", |
| elem_classes=["gradient-selector"], |
| ) |
| mode = gr.Radio( |
| ["Fast DDIM · 100 steps", "Detailed DDIM · 200 steps", "Full DDPM · 1000 steps"], |
| value="Full DDPM · 1000 steps", |
| label="Sampling path", |
| elem_classes=["gradient-selector"], |
| ) |
| curated = gr.Checkbox(value=True, label="Use visually audited FP32 seed") |
| run = gr.Button( |
| "Start diffusion", |
| variant="primary", |
| size="lg", |
| elem_id="run-button", |
| ) |
|
|
| run.click( |
| generate, |
| inputs=[digit, model, mode, seed, curated], |
| outputs=[output, status], |
| concurrency_limit=1, |
| ) |
|
|
| gr.Markdown( |
| "[Source](https://github.com/elixpo/monodiffuse) · " |
| "[Circuit](https://github.com/elixpo/monodiffuse/tree/main/hardware/diffuino_selector)", |
| elem_id="footer-note", |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| demo.queue(max_size=16).launch( |
| css=CSS, |
| theme=gr.themes.Soft(primary_hue="violet", secondary_hue="cyan"), |
| ) |
|
|
