add: orange theme and 32 lower lim

app.py

@@ -24,12 +24,17 @@ def predict(label: int, steps: int):
         yield image, f"trajectory checkpoint {idx}/{total} | integration step {step_idx}/{total_steps}"
 
 
-with gr.Blocks(title="MNIST Flow") as demo:
+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 16 trajectory checkpoints."
+        "Sampling caps ODE steps at 128 and shows 32 trajectory checkpoints."
     )
 
     grid = gr.Image(label="Trajectory", show_label=True)
@@ -37,9 +42,9 @@ with gr.Blocks(title="MNIST Flow") as demo:
 
     with gr.Row():
         label = gr.Dropdown([str(i) for i in range(10)], value="1", label="Label")
-        steps = gr.Slider(16, 128, value=128, step=1, label="Steps")
+        steps = gr.Slider(32, 128, value=128, step=1, label="Steps")
 
-    generate_btn = gr.Button("Generate")
+    generate_btn = gr.Button("Generate", variant="primary")
 
     generate_btn.click(
         fn=predict,

model.py

@@ -39,7 +39,7 @@ MODEL_CONFIG = {
 INFER_CONFIG = {
     "steps": 128,
     "cfg_scale": 2.0,
-    "trajectory_checkpoints": 16,
+    "trajectory_checkpoints": 32,
 }
 
 DTYPES = {
@@ -543,7 +543,7 @@ def flow_image_generate(
         uncond_context = uncond_context.to(device=context.device, dtype=context.dtype)
 
     history: list[tuple[int, torch.Tensor]] = []
-    checkpoint_count = max(16, int(INFER_CONFIG["trajectory_checkpoints"]))
+    checkpoint_count = max(32, int(INFER_CONFIG["trajectory_checkpoints"]))
     checkpoint_indices = np.linspace(1, steps, num=checkpoint_count, dtype=int).tolist()
     checkpoint_indices = sorted(set(max(1, min(steps, idx)) for idx in checkpoint_indices))
 
@@ -677,7 +677,7 @@ def generate_grid_image(label: int, steps: int, num_samples: int) -> Image.Image
 @torch.inference_mode()
 def iter_trajectory_frames(label: int, steps: int):
     model, device, _ = load_model()
-    steps = max(16, min(int(steps), 128))
+    steps = max(32, min(int(steps), 128))
     context = torch.full((1,), int(label), device=device, dtype=torch.long)
     prompt = torch.empty((1, 0), device=device, dtype=torch.long)
     null_label_id = int(MODEL_CONFIG["null_label_id"])