hi

app.py

@@ -1,3 +1,4 @@
+import os
 import spaces
 from pathlib import Path
 import argparse
@@ -20,11 +21,11 @@ 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
-LOOP_S = 0.3           # autoplay speed (s)
+NUM_FRAMES = 9  # frame_0.png ... frame_8.png
+
 
 @spaces.GPU(timeout=300, duration=80)
-def generate_vstack_from_image(image: Image.Image, input_focal_position: int, num_inference_steps: int):
+def generate_outputs(image: Image.Image, input_focal_position: int, num_inference_steps: int):
     if image is None:
         raise gr.Error("Please upload an image first.")
 
@@ -37,35 +38,37 @@ def generate_vstack_from_image(image: Image.Image, input_focal_position: int, nu
 
     write_output(OUTPUT_DIR, output_frames, focal_stack_num, batch["icc_profile"])
 
+    video_path = OUTPUT_DIR / "stack.mp4"
     first_frame = OUTPUT_DIR / "frame_0.png"
+
+    if not video_path.exists():
+        raise gr.Error("stack.mp4 not found in output_dir")
     if not first_frame.exists():
         raise gr.Error("frame_0.png not found in output_dir")
 
-    # Show first frame + reset slider
-    return str(first_frame), gr.update(value=0)
+    return str(video_path), str(first_frame), gr.update(value=0)
+
 
 def show_frame(idx: int):
-    path = OUTPUT_DIR / f"frame_{idx}.png"
+    path = OUTPUT_DIR / f"frame_{int(idx)}.png"
     if not path.exists():
         return None
     return str(path)
 
-def advance_frame(idx: int, autoplay: bool):
-    if not autoplay:
-        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():
-        return gr.update(value=int(idx)), gr.update()
+def set_view_mode(mode: str):
+    show_video = (mode == "Video")
+    return (
+        gr.update(visible=show_video),
+        gr.update(visible=not show_video),
+    )
 
-    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. Toggle autoplay on the output panel.
+        Switch between **Video** (stack.mp4) and **Frames** (PNG slider).
         """
     )
 
@@ -79,6 +82,7 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                 maximum=8,
                 step=1,
                 value=4,
+                interactive=True,
             )
 
             num_inference_steps = gr.Slider(
@@ -93,22 +97,36 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
             generate_btn = gr.Button("Generate stack", variant="primary")
 
         with gr.Column():
-            frame_view = gr.Image(label="Frame viewer", type="filepath")
+            view_mode = gr.Radio(
+                choices=["Video", "Frames"],
+                value="Video",
+                label="Output view",
+            )
 
-            frame_slider = gr.Slider(
-                minimum=0,
-                maximum=NUM_FRAMES - 1,
-                step=1,
-                value=0,
-                label="Output Focus Position",
+            # --- Video output ---
+            video_out = gr.Video(
+                label="Generated stack (stack.mp4)",
+                format="mp4",
+                autoplay=True,
+                loop=True,
+                visible=True,
             )
 
-            autoplay = gr.Checkbox(value=True, label="Autoplay")
+            # --- Frames output (group) ---
+            with gr.Group(visible=False) as frames_group:
+                frame_view = gr.Image(label="Frame viewer", type="filepath")
+                frame_slider = gr.Slider(
+                    minimum=0,
+                    maximum=NUM_FRAMES - 1,
+                    step=1,
+                    value=0,
+                    label="Output focal position",
+                )
 
     generate_btn.click(
-        fn=generate_vstack_from_image,
+        fn=generate_outputs,
         inputs=[image_in, input_focal_position, num_inference_steps],
-        outputs=[frame_view, frame_slider],
+        outputs=[video_out, frame_view, frame_slider],
         api_name="predict",
     )
 
@@ -118,11 +136,10 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
         outputs=frame_view,
     )
 
-    timer = gr.Timer(LOOP_S)
-    timer.tick(
-        fn=advance_frame,
-        inputs=[frame_slider, autoplay],
-        outputs=[frame_slider, frame_view],
+    view_mode.change(
+        fn=set_view_mode,
+        inputs=view_mode,
+        outputs=[video_out, frames_group],
     )
 
 if __name__ == "__main__":