app.py

import torch
import gradio as gr
import spaces

from model import iter_trajectory_frames, load_model

MODEL_READY = False


def ensure_model_loaded():
    global MODEL_READY
    if not MODEL_READY:
        load_model()
        MODEL_READY = True


@spaces.GPU
@torch.inference_mode()
def predict(label: int, steps: int):
    ensure_model_loaded()
    for idx, (image, step_idx, total_steps, total) in enumerate(
        iter_trajectory_frames(label=label, steps=steps), start=1
    ):
        yield image, f"trajectory checkpoint {idx}/{total} | denoising step {step_idx}/{total_steps}"


with gr.Blocks(title="MNIST Diffusion") as demo:
    gr.Markdown("# MNIST Diffusion")
    gr.Markdown(
        "Discrete diffusion model for MNIST digits. "
        "The demo streams one sample as masked tokens are resolved with fixed CFG=2.0, "
        "temperature=0.6, and top_p=0.99."
    )

    grid = gr.Image(label="Trajectory", show_label=True)
    status = gr.Textbox(label="Status")

    with gr.Row():
        label = gr.Dropdown([str(i) for i in range(10)], value="4", label="Label")
        steps = gr.Slider(32, 784, value=784, step=1, label="Steps")

    generate_btn = gr.Button("Generate")

    generate_btn.click(
        fn=predict,
        inputs=[label, steps],
        outputs=[grid, status],
        scroll_to_output=True,
    )

if __name__ == "__main__":
    demo.launch()