try again

app.py

@@ -1,6 +1,5 @@
 import os
 import spaces
-import uuid
 from pathlib import Path
 import argparse
 
@@ -22,82 +21,53 @@ 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
+
 @spaces.GPU(timeout=300, duration=80)
-def generate_vstack_from_image(
-    image: Image.Image,
-    input_focal_position: int,
-    num_inference_steps: int,
-):
-    """
-    Wrapper for Gradio.
-    Generates a focal stack and returns:
-      - list of PNG frame paths (state)
-      - first frame for display
-      - slider configuration
-    """
+def generate_vstack_from_image(image: Image.Image, input_focal_position: int, num_inference_steps: int):
     if image is None:
         raise gr.Error("Please upload an image first.")
 
     args.num_inference_steps = num_inference_steps
     args.device = "cuda"
-
     pipe.to(args.device)
 
     batch = convert_to_batch(image, input_focal_position=input_focal_position)
     output_frames, focal_stack_num = inference_on_image(args, batch, pipe, device)
 
-    run_id = uuid.uuid4().hex
-    save_dir = OUTPUT_DIR / run_id
-    save_dir.mkdir(parents=True, exist_ok=True)
-
-    write_output(save_dir, output_frames, focal_stack_num, batch["icc_profile"])
-
-    # ---- SIMPLE FIND: PNGs only, sorted ----
-    frame_paths = sorted(
-        Path(save_dir).glob("*.png"),
-        key=lambda p: p.stem,
+    write_output(
+        OUTPUT_DIR,
+        output_frames,
+        focal_stack_num,
+        batch["icc_profile"],
     )
 
-    if len(frame_paths) == 0:
-        raise gr.Error("No PNG frames found in output directory.")
-
-    slider = gr.Slider(
-        minimum=0,
-        maximum=len(frame_paths) - 1,
-        step=1,
-        value=0,
-        label="Frame index",
-    )
+    # 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(p) for p in frame_paths], str(frame_paths[0]), slider
+    return str(first_frame)
 
 
-def show_frame(frame_paths, idx: int):
-    if not frame_paths:
+def show_frame(idx: int):
+    path = OUTPUT_DIR / f"frame_{idx}.png"
+    if not path.exists():
         return None
-    idx = max(0, min(int(idx), len(frame_paths) - 1))
-    return frame_paths[idx]
+    return str(path)
 
 
 with gr.Blocks(css="footer {visibility: hidden}") as demo:
     gr.Markdown(
         """
         # 🖼️ ➜ 🎬 Generate Focal Stacks from a Single Image
-
-        Upload an image, set the input focal position, and generate a focal stack.
-        Use the slider to scrub through the saved frames.
+        Generate a focal stack and scrub through frames using the slider.
         """
     )
 
-    frame_paths_state = gr.State([])
-
     with gr.Row():
         with gr.Column():
-            image_in = gr.Image(
-                type="pil",
-                label="Input image",
-                interactive=True,
-            )
+            image_in = gr.Image(type="pil", label="Input image", interactive=True)
 
             input_focal_position = gr.Slider(
                 label="Input focal position (Near - 5cm, Far - Infinity):",
@@ -105,7 +75,6 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                 maximum=8,
                 step=1,
                 value=4,
-                interactive=True,
             )
 
             num_inference_steps = gr.Slider(
@@ -114,34 +83,30 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                 maximum=25,
                 step=1,
                 value=25,
-                info="More steps = better quality but slower",
             )
 
             generate_btn = gr.Button("Generate stack", variant="primary")
 
         with gr.Column():
-            frame_view = gr.Image(
-                label="Frame viewer",
-                type="filepath",
-            )
+            frame_view = gr.Image(label="Frame viewer", type="filepath")
+
             frame_slider = gr.Slider(
                 minimum=0,
-                maximum=0,
+                maximum=NUM_FRAMES - 1,
                 step=1,
                 value=0,
-                label="Frame index",
+                label="Focal plane",
             )
 
     generate_btn.click(
         fn=generate_vstack_from_image,
         inputs=[image_in, input_focal_position, num_inference_steps],
-        outputs=[frame_paths_state, frame_view, frame_slider],
-        api_name="predict",
+        outputs=frame_view,
     )
 
     frame_slider.change(
         fn=show_frame,
-        inputs=[frame_paths_state, frame_slider],
+        inputs=frame_slider,
         outputs=frame_view,
     )