import gradio as gr
import spaces
import torch

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} | integration step {step_idx}/{total_steps}"


theme = gr.themes.Default(
    primary_hue="orange",
)


with gr.Blocks(title="MNIST Flow", theme=theme) as demo:
    gr.Markdown("# MNIST Flow")
    gr.Markdown(
        "Flow-matching DiT model for MNIST digits. "
        "The demo streams one sample from noise to digit with fixed CFG=2.0. "
        "Sampling caps ODE steps at 128 and shows 32 trajectory checkpoints."
    )

    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="1", label="Label")
        steps = gr.Slider(32, 128, value=128, step=1, label="Steps")

    generate_btn = gr.Button("Generate", variant="primary")

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


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