autoplay:

app.py

@@ -1,4 +1,3 @@
-import os
 import spaces
 from pathlib import Path
 import argparse
@@ -21,7 +20,8 @@ pipe, device = load_model(args)
 OUTPUT_DIR = Path("/tmp/output_stacks")
 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 
-NUM_FRAMES = 9  # frame_0.png ... frame_8.png
+NUM_FRAMES = 9          # frame_0.png ... frame_8.png
+LOOP_MS = 300           # autoplay speed (ms)
 
 @spaces.GPU(timeout=300, duration=80)
 def generate_vstack_from_image(image: Image.Image, input_focal_position: int, num_inference_steps: int):
@@ -35,20 +35,14 @@ def generate_vstack_from_image(image: Image.Image, input_focal_position: int, nu
     batch = convert_to_batch(image, input_focal_position=input_focal_position)
     output_frames, focal_stack_num = inference_on_image(args, batch, pipe, device)
 
-    write_output(
-        OUTPUT_DIR,
-        output_frames,
-        focal_stack_num,
-        batch["icc_profile"],
-    )
+    write_output(OUTPUT_DIR, output_frames, focal_stack_num, batch["icc_profile"])
 
-    # Show first frame immediately
     first_frame = OUTPUT_DIR / "frame_0.png"
     if not first_frame.exists():
         raise gr.Error("frame_0.png not found in output_dir")
 
-    return str(first_frame)
-
+    # show first frame + reset slider to 0
+    return str(first_frame), gr.update(value=0)
 
 def show_frame(idx: int):
     path = OUTPUT_DIR / f"frame_{idx}.png"
@@ -56,12 +50,25 @@ def show_frame(idx: int):
         return None
     return str(path)
 
+def advance_frame(idx: int, autoplay: bool):
+    # If autoplay is off, don't change anything.
+    if not autoplay:
+        # keep slider value, keep image (returning None keeps the current image)
+        return gr.update(value=int(idx)), gr.update()
+
+    next_idx = (int(idx) + 1) % NUM_FRAMES
+    next_path = OUTPUT_DIR / f"frame_{next_idx}.png"
+    if not next_path.exists():
+        # If frames aren't ready yet, hold.
+        return gr.update(value=int(idx)), gr.update()
+
+    return gr.update(value=next_idx), str(next_path)
 
 with gr.Blocks(css="footer {visibility: hidden}") as demo:
     gr.Markdown(
         """
         # 🖼️ ➜ 🎬 Generate Focal Stacks from a Single Image
-        Generate a focal stack and scrub through frames using the slider.
+        Generate a focal stack and scrub through frames using the slider. Toggle autoplay to loop.
         """
     )
 
@@ -83,8 +90,11 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                 maximum=25,
                 step=1,
                 value=25,
+                info="More steps = better quality but slower",
             )
 
+            autoplay = gr.Checkbox(value=True, label="Autoplay")
+
             generate_btn = gr.Button("Generate stack", variant="primary")
 
         with gr.Column():
@@ -98,17 +108,28 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                 label="Focal plane",
             )
 
+    # Generate outputs + reset slider
     generate_btn.click(
         fn=generate_vstack_from_image,
         inputs=[image_in, input_focal_position, num_inference_steps],
-        outputs=frame_view,
+        outputs=[frame_view, frame_slider],
+        api_name="predict",
     )
 
+    # Manual scrubbing
     frame_slider.change(
         fn=show_frame,
         inputs=frame_slider,
         outputs=frame_view,
     )
 
+    # Autoplay loop (can be toggled off via checkbox)
+    timer = gr.Timer(LOOP_MS)
+    timer.tick(
+        fn=advance_frame,
+        inputs=[frame_slider, autoplay],
+        outputs=[frame_slider, frame_view],
+    )
+
 if __name__ == "__main__":
     demo.launch()