Update app.py

app.py

@@ -1,15 +1,16 @@
 # app.py
-# A Gradio Space that (1) extracts frames with FFmpeg and (2) upscales frames with Real-ESRGAN.
-# NEW: You can now also upload existing images directly to the Upscale tab (no video needed).
-# - CPU-friendly (GPU optional)
-# - Exact ffmpeg command preview
-# - ZIP download for results
+# Two-step UI:
+#  1) Extract frames (Video → JPG/PNG)
+#  2) After extraction completes, the Upscale section becomes visible and can be run separately
+# Extras:
+#  - Preview shows 30 frames sampled evenly across the whole video (scrollable gallery)
+#  - Prefix defaults to the input video filename if left blank
+#  - Separate tab for directly upscaling uploaded images (optional)
 
 import os
-import io
 import re
-import math
 import json
+import math
 import time
 import shutil
 import zipfile
@@ -23,7 +24,7 @@ import numpy as np
 from PIL import Image
 
 # ─────────────────────────────────────────────────────────────
-# FFmpeg detection
+# System checks
 # ─────────────────────────────────────────────────────────────
 
 def _which(name: str) -> Optional[str]:
@@ -35,15 +36,38 @@ FFPROBE = _which("ffprobe")
 
 if not FFMPEG or not FFPROBE:
     MISSING_MSG = (
-        "⚠️ FFmpeg not found. Add a 'packages.txt' with just: \nffmpeg\nlibsm6\nlibxext6\nThen restart the Space."
+        "⚠️ FFmpeg not found. Add a 'packages.txt' with exactly:
+ffmpeg
+libsm6
+libxext6
+Then restart the Space."
     )
 else:
     MISSING_MSG = ""
 
+# Real-ESRGAN (optional but needed for step 2)
+try:
+    from realesrgan import RealESRGANer
+    from basicsr.archs.rrdbnet_arch import RRDBNet
+    _HAVE_REALESRGAN = True
+except Exception:
+    _HAVE_REALESRGAN = False
+
 # ─────────────────────────────────────────────────────────────
-# ffprobe helpers
+# Helpers
 # ─────────────────────────────────────────────────────────────
 
+def sample_paths(paths, n=30):
+    """Return up to n items sampled evenly across the list, preserving order."""
+    if not paths:
+        return []
+    n = max(1, min(n, len(paths)))
+    idxs = np.linspace(0, len(paths) - 1, num=n, dtype=int).tolist()
+    # Ensure uniq & sorted
+    idxs = sorted(dict.fromkeys(idxs))
+    return [paths[i] for i in idxs]
+
+
 def ffprobe_json(input_path: str) -> dict:
     if not FFPROBE:
         return {}
@@ -82,9 +106,6 @@ def parse_video_info(meta: dict) -> dict:
         info["height"] = v.get("height")
     return info
 
-# ─────────────────────────────────────────────────────────────
-# FFmpeg command builder
-# ─────────────────────────────────────────────────────────────
 
 def build_ffmpeg_command(
     input_path: str,
@@ -143,21 +164,10 @@ def build_ffmpeg_command(
     cmd += ["-frame_pts", "1", out_pattern]
     return cmd
 
-# ─────────────────────────────────────────────────────────────
-# Real-ESRGAN setup (CPU/GPU)
-# ─────────────────────────────────────────────────────────────
-try:
-    from realesrgan import RealESRGANer
-    from basicsr.archs.rrdbnet_arch import RRDBNet
-    _HAVE_REALESRGAN = True
-except Exception:
-    _HAVE_REALESRGAN = False
-
 
 def get_realesrganer(model_name: str, scale: int, tile: int, half: bool, device: str = "cpu"):
     if not _HAVE_REALESRGAN:
         raise RuntimeError("realesrgan is not installed. Check requirements.txt")
-    # Model selection
     if model_name in ("x4plus", "x4plus-anime"):
         model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
         model_scale = 4
@@ -172,7 +182,7 @@ def get_realesrganer(model_name: str, scale: int, tile: int, half: bool, device:
 
     upsampler = RealESRGANer(
         scale=model_scale,
-        model_path=None,  # let library fetch weights
+        model_path=None,
         model=model,
         tile=tile or 0,
         tile_pad=10,
@@ -180,13 +190,15 @@ def get_realesrganer(model_name: str, scale: int, tile: int, half: bool, device:
         half=half,
         device=device,
     )
-    return upsampler, model_scale
+    return upsampler
 
 
 def upscale_images(img_paths: List[Path], out_dir: Path, model_name: str, scale: int, tile: int, precision: str, progress=gr.Progress(track_tqdm=True)) -> List[str]:
+    if not _HAVE_REALESRGAN:
+        raise RuntimeError("realesrgan not available")
     device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
     half = (precision == "half") and (device == "cuda")
-    upsampler, model_scale = get_realesrganer(model_name, scale, tile, half, device=device)
+    upsampler = get_realesrganer(model_name, scale, tile, half, device=device)
 
     out_paths: List[str] = []
     for i, p in enumerate(img_paths, 1):
@@ -194,10 +206,7 @@ def upscale_images(img_paths: List[Path], out_dir: Path, model_name: str, scale:
             img = Image.open(p).convert("RGB")
             output, _ = upsampler.enhance(np.array(img), outscale=scale)
             out_img = Image.fromarray(output)
-            out_file = out_dir / p.name
-            # Save as JPG by extension; keep name but force .jpg to be consistent
-            if out_file.suffix.lower() not in [".jpg", ".jpeg"]:
-                out_file = out_file.with_suffix(".jpg")
+            out_file = out_dir / (p.stem + ".jpg")
             out_img.save(out_file, quality=95)
             out_paths.append(str(out_file))
         except Exception as e:
@@ -209,7 +218,7 @@ def upscale_images(img_paths: List[Path], out_dir: Path, model_name: str, scale:
 # Pipelines
 # ─────────────────────────────────────────────────────────────
 
-def run_video_pipeline(
+def run_video_extract(
     video: gr.File | None,
     mode: str,
     every_seconds: float,
@@ -224,25 +233,16 @@ def run_video_pipeline(
     scene_detect: bool,
     scene_thresh: float,
     prefix: str,
-    do_upscale: bool,
-    model_name: str,
-    scale: int,
-    tile: int,
-    precision: str,
 ):
     if not video or not video.name:
-        return None, None, None, "Upload a video.", ""
+        return None, None, "Upload a video.", "", gr.update(visible=False), [], "", ""
     if not FFMPEG or not FFPROBE:
-        return None, None, None, "FFmpeg missing. See note below.", MISSING_MSG
-    if do_upscale and not _HAVE_REALESRGAN:
-        return None, None, None, "realesrgan is not installed (see requirements.txt)", ""
+        return None, None, "FFmpeg missing. See note below.", MISSING_MSG, gr.update(visible=False), [], "", ""
 
     # Work dirs
-    work = Path(tempfile.mkdtemp(prefix="vid2up_"))
+    work = Path(tempfile.mkdtemp(prefix="vid2img_"))
     raw_dir = work / "frames_raw"
-    up_dir = work / "frames_upscaled"
     raw_dir.mkdir(parents=True, exist_ok=True)
-    up_dir.mkdir(parents=True, exist_ok=True)
 
     # Default prefix from input filename if blank
     if not prefix or not prefix.strip():
@@ -266,35 +266,62 @@ def run_video_pipeline(
         scene_thresh=scene_thresh,
         out_pattern=pattern,
     )
-
     cmd_preview = " ".join([s if " " not in s else f'"{s}"' for s in cmd])
 
     proc = subprocess.run(cmd, capture_output=True, text=True)
     if proc.returncode != 0:
-        return None, None, None, f"FFmpeg error:\n{proc.stderr}", cmd_preview
+        return None, None, f"FFmpeg error:
+{proc.stderr}", cmd_preview, gr.update(visible=False), [], "", ""
 
     frames = sorted(raw_dir.glob(f"{prefix}_*.{out_format}"))
     if not frames:
-        return None, None, None, "No frames extracted.", cmd_preview
-
-    # Optionally upscale
-    if do_upscale:
-        up_paths = upscale_images(frames, up_dir, model_name, scale, tile, precision)
-        gallery = up_paths[:60]
-        final_zip_dir = up_dir
-        detail = f"Frames extracted: {len(frames)} | Upscaled: {len(up_paths)} | Model: {model_name} | Scale: x{scale} | Tile: {tile} | Precision: {precision}"
-    else:
-        gallery = [str(p) for p in frames[:60]]
-        final_zip_dir = raw_dir
-        detail = f"Frames extracted: {len(frames)} | (Upscale: off)"
+        return None, None, "No frames extracted.", cmd_preview, gr.update(visible=False), [], "", ""
+
+    # Preview (30 sampled)
+    gallery = [str(p) for p in sample_paths(frames, n=30)]
 
     # Zip
-    zip_path = work / ("upscaled_frames.zip" if do_upscale else "frames.zip")
+    zip_path = work / "frames.zip"
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for p in frames:
+            zf.write(p, p.name)
+
+    details = f"Frames extracted: {len(frames)} | Saved to: {raw_dir}"
+
+    # Make upscale section visible and return state for next step
+    return gallery, str(zip_path), details, cmd_preview, gr.update(visible=True), [str(p) for p in frames], str(raw_dir), prefix
+
+
+def run_upscale_from_extracted(
+    frames_list: List[str] | None,
+    frames_dir: str,
+    prefix: str,
+    model_name: str,
+    scale: int,
+    tile: int,
+    precision: str,
+):
+    if not _HAVE_REALESRGAN:
+        return None, None, "realesrgan is not installed (see requirements.txt)", ""
+    if not frames_list:
+        return None, None, "No extracted frames state found. Please run extraction first.", ""
+
+    work = Path(tempfile.mkdtemp(prefix="up_"))
+    out_dir = work / "upscaled"
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    img_paths = [Path(p) for p in frames_list]
+    up_paths = upscale_images(img_paths, out_dir, model_name, scale, tile, precision)
+
+    gallery = sample_paths([Path(p) for p in up_paths], n=30)
+
+    zip_path = work / "upscaled_frames.zip"
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
-        for p in sorted(final_zip_dir.glob("*")):
+        for p in sorted(out_dir.glob("*")):
             zf.write(p, p.name)
 
-    return gallery, str(zip_path), detail, "", cmd_preview
+    detail = f"Upscaled: {len(up_paths)} | Model: {model_name} | Scale: x{scale} | Tile: {tile} | Precision: {precision}"
+    return gallery, str(zip_path), detail, ""
 
 
 def run_image_upscale_pipeline(
@@ -316,20 +343,18 @@ def run_image_upscale_pipeline(
     in_dir.mkdir(parents=True, exist_ok=True)
     out_dir.mkdir(parents=True, exist_ok=True)
 
-    # Save uploads into a temp folder with cleaned names
     img_paths: List[Path] = []
     for f in images:
         p = Path(f.name)
         name = p.name
         if prefix and prefix.strip():
-            # keep original number if present else add incremental index
             name = f"{prefix}_{name}"
         dst = in_dir / name
         shutil.copy2(p, dst)
         img_paths.append(dst)
 
     up_paths = upscale_images(img_paths, out_dir, model_name, scale, tile, precision)
-    gallery = up_paths[:60]
+    gallery = [str(p) for p in sample_paths([Path(p) for p in up_paths], n=30)]
 
     zip_path = work / "upscaled_images.zip"
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
@@ -349,13 +374,18 @@ def build_ui():
     .cmdbox textarea { font-family: ui-monospace, Menlo, monospace; font-size: 12px; }
     """) as demo:
         gr.Markdown("""
-        <div class="cf-title">Video → JPG → Upscale (FFmpeg + Real-ESRGAN)</div>
-        Extract frames from video or upload images directly, then upscale with Real-ESRGAN.
+        <div class=\"cf-title\">Video → JPG → Upscale (Two-step)</div>
+        Step 1 extracts frames. When it finishes, Step 2 (Upscale) appears.
         """)
 
+        # Shared state between steps
+        frames_state = gr.State([])     # list[str] of extracted frame paths
+        frames_dir_state = gr.State("")
+        prefix_state = gr.State("")
+
         with gr.Tabs():
-            # Tab 1: Video → Frames → (Upscale)
-            with gr.Tab("Video pipeline"):
+            # ── Tab 1: Video → Frames (Step 1)
+            with gr.Tab("Video → Frames"):
                 with gr.Row():
                     video = gr.File(label="Upload video", file_types=[".mp4", ".mov", ".mkv", ".avi", ".webm", ".m4v"], type="filepath")
                 with gr.Accordion("Extraction", open=True):
@@ -376,18 +406,10 @@ def build_ui():
                         scene_detect = gr.Checkbox(False, label="Scene-change detect")
                         scene_thresh = gr.Slider(0.0, 1.0, value=0.3, step=0.01, label="Scene threshold")
                         prefix_vid = gr.Textbox(value="", label="Filename prefix (defaults to input file name)")
-                with gr.Accordion("Upscaling", open=True):
-                    with gr.Row():
-                        do_upscale = gr.Checkbox(True, label="Upscale frames with Real-ESRGAN")
-                        model_name = gr.Dropdown(["x4plus", "x4plus-anime", "x2plus"], value="x4plus", label="Model")
-                        scale = gr.Dropdown([2, 4], value=4, label="Output scale")
-                    with gr.Row():
-                        tile = gr.Number(value=0, label="Tile size (0 = auto)")  # tiling reduces RAM; try 256/512 on GPU, 0 for auto
-                        precision = gr.Dropdown(["auto", "half", "full"], value="auto", label="Precision (GPU=half, CPU=full)")
 
-                run_btn = gr.Button("Run: Extract → (Upscale) → ZIP", variant="primary")
-                gallery = gr.Gallery(label="Preview (first 60)", columns=6, height=300)
-                zip_out = gr.File(label="Download ZIP")
+                run_extract_btn = gr.Button("Step 1: Extract frames", variant="primary")
+                gallery = gr.Gallery(label="Preview (30 sampled)", columns=6, height=480)
+                zip_out = gr.File(label="Download frames ZIP")
                 details = gr.Markdown("Ready.")
                 with gr.Accordion("Show FFmpeg command", open=False):
                     cmd_preview = gr.Textbox(label="ffmpeg command", lines=4, elem_classes=["cmdbox"])
@@ -407,39 +429,58 @@ def build_ui():
                 out_format.change(_toggle, [mode, out_format], [every_seconds, nth_frame, exact_fps, jpg_quality, png_level])
                 demo.load(_toggle, [mode, out_format], [every_seconds, nth_frame, exact_fps, jpg_quality, png_level])
 
-                run_btn.click(
-                    run_video_pipeline,
+                # Run extraction → also reveal upscale group and set states
+                show_upscale_group = gr.State(False)
+                upscale_group = gr.Group(visible=False)
+
+                run_extract_btn.click(
+                    run_video_extract,
                     inputs=[
                         video, mode, every_seconds, nth_frame, exact_fps,
                         start_time, end_time, long_side, out_format, jpg_quality, png_level,
                         scene_detect, scene_thresh, prefix_vid,
-                        do_upscale, model_name, scale, tile, precision,
                     ],
-                    outputs=[gallery, zip_out, details, gr.Textbox(), cmd_preview],
-                    api_name="extract_and_upscale",
+                    outputs=[gallery, zip_out, details, cmd_preview, upscale_group, frames_state, frames_dir_state, prefix_state],
                 )
 
-            # Tab 2: Upscale images directly
-            with gr.Tab("Upscale images"):
-                with gr.Row():
-                    imgs = gr.Files(label="Upload images (JPG/PNG)", file_types=[".jpg", ".jpeg", ".png"], type="filepath")
+            # ── Tab 2: Upscale (Step 2 — appears after extraction)
+            with gr.Tab("Upscale extracted", visible=True):
+                with gr.Group() as upscale_group:  # rebind same name for clarity
+                    gr.Markdown("**Step 2:** Upscale the frames you just extracted.")
+                    with gr.Row():
+                        model_name = gr.Dropdown(["x4plus", "x4plus-anime", "x2plus"], value="x4plus", label="Model")
+                        scale = gr.Dropdown([2, 4], value=4, label="Output scale")
+                        tile = gr.Number(value=0, label="Tile size (0 = auto)")
+                        precision = gr.Dropdown(["auto", "half", "full"], value="auto", label="Precision (GPU=half, CPU=full)")
+                    run_upscale_btn = gr.Button("Step 2: Upscale extracted frames", variant="primary")
+                    gallery_up = gr.Gallery(label="Upscaled preview (30 sampled)", columns=6, height=480)
+                    zip_up = gr.File(label="Download upscaled ZIP")
+                    details_up = gr.Markdown("")
+
+                    run_upscale_btn.click(
+                        run_upscale_from_extracted,
+                        inputs=[frames_state, frames_dir_state, prefix_state, model_name, scale, tile, precision],
+                        outputs=[gallery_up, zip_up, details_up, gr.Textbox()],
+                    )
+
+            # ── Optional: direct image upscale
+            with gr.Tab("Upscale images (upload)"):
+                imgs = gr.Files(label="Upload images (JPG/PNG)", file_types=[".jpg", ".jpeg", ".png"], type="filepath")
                 with gr.Row():
                     model_name_i = gr.Dropdown(["x4plus", "x4plus-anime", "x2plus"], value="x4plus", label="Model")
                     scale_i = gr.Dropdown([2, 4], value=4, label="Output scale")
                     tile_i = gr.Number(value=0, label="Tile size (0 = auto)")
                     precision_i = gr.Dropdown(["auto", "half", "full"], value="auto", label="Precision (GPU=half, CPU=full)")
-                with gr.Row():
-                    prefix_img = gr.Textbox(value="", label="Optional filename prefix (adds prefix_ to each output)")
-                run_btn_i = gr.Button("Run: Upscale → ZIP", variant="primary")
-                gallery_i = gr.Gallery(label="Preview (first 60)", columns=6, height=300)
+                prefix_img = gr.Textbox(value="", label="Optional filename prefix (adds prefix_ to each output)")
+                run_btn_i = gr.Button("Upscale uploaded images", variant="secondary")
+                gallery_i = gr.Gallery(label="Preview (30 sampled)", columns=6, height=480)
                 zip_out_i = gr.File(label="Download ZIP")
-                details_i = gr.Markdown("Ready.")
+                details_i = gr.Markdown("")
 
                 run_btn_i.click(
                     run_image_upscale_pipeline,
                     inputs=[imgs, model_name_i, scale_i, tile_i, precision_i, prefix_img],
                     outputs=[gallery_i, zip_out_i, details_i, gr.Textbox()],
-                    api_name="upscale_images_only",
                 )
 
     return demo