Update app.py

app.py

@@ -1,11 +1,14 @@
+# =============================
 # app.py
-# 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)
+# FFmpeg Frames + Real-ESRGAN Upscale + Re-encode (3-step) + Quick Mode
+# - Step 1: Extract frames (with live estimate & progress, Cancel)
+# - Step 2: Upscale extracted frames (progress, Cancel)
+# - Step 3: Re-encode frames to MP4/WebM with optional original audio (progress, Cancel)
+# - Quick Mode: One-click pipeline (All Frames → Upscale x4 → MP4 H.264 with audio)
+# - Previews show 30 frames sampled across the whole set; galleries are scrollable
+# - Prefix defaults to input video filename if left blank
+# - Direct image-upscale tab also included
+# =============================
 
 import os
 import re
@@ -17,14 +20,14 @@ import zipfile
 import tempfile
 import subprocess
 from pathlib import Path
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import gradio as gr
 import numpy as np
 from PIL import Image
 
 # ─────────────────────────────────────────────────────────────
-# System checks
+# System checks & deps
 # ─────────────────────────────────────────────────────────────
 
 def _which(name: str) -> Optional[str]:
@@ -34,27 +37,39 @@ def _which(name: str) -> Optional[str]:
 FFMPEG = _which("ffmpeg")
 FFPROBE = _which("ffprobe")
 
-# Real-ESRGAN (optional but needed for step 2)
+if not FFMPEG or not FFPROBE:
+    MISSING_MSG = (
+        "⚠️ FFmpeg not found. Add a 'packages.txt' with exactly:\nffmpeg\nlibsm6\nlibxext6\nThen restart the Space."
+    )
+else:
+    MISSING_MSG = ""
+
 try:
     from realesrgan import RealESRGANer
     from basicsr.archs.rrdbnet_arch import RRDBNet
-    _HAVE_REALESRGAN = True
+    HAVE_REALESRGAN = True
 except Exception:
-    _HAVE_REALESRGAN = False
+    HAVE_REALESRGAN = False
 
 # ─────────────────────────────────────────────────────────────
 # Helpers
 # ─────────────────────────────────────────────────────────────
 
-def sample_paths(paths, n=30):
-    """Return up to n items sampled evenly across the list, preserving order."""
+def sample_paths(paths: List[Path] | List[str], n: int = 30) -> List[str]:
+    """Return up to n items sampled evenly across the list, preserving order (as strings)."""
     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]
+    return [str(paths[i]) for i in idxs]
+
+
+def sanitize_prefix(txt: str) -> str:
+    txt = (txt or "").strip()
+    if not txt:
+        return ""
+    return re.sub(r"[^A-Za-z0-9._-]+", "_", txt)[:80]
 
 
 def ffprobe_json(input_path: str) -> dict:
@@ -83,10 +98,9 @@ def parse_video_info(meta: dict) -> dict:
         v = vstreams[0]
         rfr = v.get("r_frame_rate") or v.get("avg_frame_rate")
         if rfr and "/" in rfr:
-            num, den = rfr.split("/")
             try:
-                num = float(num)
-                den = float(den)
+                num, den = rfr.split("/")
+                num = float(num); den = float(den)
                 if den != 0:
                     info["fps"] = num / den
             except Exception:
@@ -96,7 +110,26 @@ def parse_video_info(meta: dict) -> dict:
     return info
 
 
-def build_ffmpeg_command(
+def estimate_output_count(mode: str, duration: float | None, in_fps: float | None,
+                           every_seconds: float, nth_frame: int, exact_fps: float) -> Optional[int]:
+    if not duration:
+        return None
+    in_fps = in_fps or 30.0
+    try:
+        if mode == "All frames":
+            return int(math.ceil(duration * in_fps))
+        if mode == "Every N seconds" and every_seconds > 0:
+            return int(math.ceil(duration / every_seconds))
+        if mode == "Every Nth frame" and nth_frame > 0:
+            return int(math.ceil((duration * in_fps) / nth_frame))
+        if mode == "Exact FPS" and exact_fps > 0:
+            return int(math.ceil(duration * exact_fps))
+    except Exception:
+        return None
+    return None
+
+
+def build_ffmpeg_extract(
     input_path: str,
     mode: str,
     every_seconds: float,
@@ -114,14 +147,12 @@ def build_ffmpeg_command(
 ) -> List[str]:
     if not FFMPEG:
         raise RuntimeError("FFmpeg not available")
-
     cmd = [FFMPEG, "-y"]
     if start_time:
         cmd += ["-ss", start_time]
     cmd += ["-i", input_path]
     if end_time:
         cmd += ["-to", end_time]
-
     vf = []
     if mode == "Every N seconds":
         vf.append(f"fps={max(1e-6, 1.0/float(every_seconds or 1))}")
@@ -134,29 +165,24 @@ def build_ffmpeg_command(
         pass
     else:
         vf.append("fps=1")
-
     if scene_detect:
         vf.append(f"select='gt(scene,{float(scene_thresh)})',showinfo")
         vf.append("setpts=N/FRAME_RATE/TB")
-
     if long_side and long_side > 0:
         vf.append("scale='if(gt(iw,ih),%d,-1)':'if(gt(iw,ih),-1,%d)':force_original_aspect_ratio=decrease" % (long_side, long_side))
-
     if vf:
         cmd += ["-vf", ",".join(vf)]
-
     if out_format == "jpg":
         cmd += ["-q:v", str(jpg_quality)]
     elif out_format == "png":
         cmd += ["-compression_level", str(png_level)]
-
     cmd += ["-frame_pts", "1", out_pattern]
     return cmd
 
 
 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")
+    if not HAVE_REALESRGAN:
+        raise RuntimeError("realesrgan is not installed. See requirements.txt")
     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
@@ -165,10 +191,8 @@ def get_realesrganer(model_name: str, scale: int, tile: int, half: bool, device:
         model_scale = 2
     else:
         raise ValueError("Unknown Real-ESRGAN model")
-
     if scale not in (2, 4):
         scale = model_scale
-
     upsampler = RealESRGANer(
         scale=model_scale,
         model_path=None,
@@ -181,33 +205,21 @@ def get_realesrganer(model_name: str, scale: int, tile: int, half: bool, device:
     )
     return upsampler
 
+# ─────────────────────────────────────────────────────────────
+# Progress UI helper (HTML bar)
+# ─────────────────────────────────────────────────────────────
 
-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 = get_realesrganer(model_name, scale, tile, half, device=device)
-
-    out_paths: List[str] = []
-    for i, p in enumerate(img_paths, 1):
-        try:
-            img = Image.open(p).convert("RGB")
-            output, _ = upsampler.enhance(np.array(img), outscale=scale)
-            out_img = Image.fromarray(output)
-            out_file = out_dir / (p.stem + ".jpg")
-            out_img.save(out_file, quality=95)
-            out_paths.append(str(out_file))
-        except Exception as e:
-            print(f"Upscale failed for {p}: {e}")
-        progress(i / max(1, len(img_paths)))
-    return out_paths
+def render_progress(pct: float, label: str = "") -> str:
+    pct = max(0.0, min(100.0, pct))
+    return f'''<div style="width:100%;border:1px solid #ddd;border-radius:8px;overflow:hidden;height:18px;">
+  <div style="height:100%;width:{pct:.1f}%;background:#3b82f6;"></div>
+</div><div style="font-size:12px;opacity:.8;margin-top:4px;">{label} {pct:.1f}%</div>'''
 
 # ─────────────────────────────────────────────────────────────
-# Pipelines
+# Step 1: Extract frames (with Cancel)
 # ─────────────────────────────────────────────────────────────
 
-def run_video_extract(
+def step1_extract(
     video: gr.File | None,
     mode: str,
     every_seconds: float,
@@ -221,32 +233,37 @@ def run_video_extract(
     png_level: int,
     scene_detect: bool,
     scene_thresh: float,
-    prefix: str,
+    prefix_in: str,
+    # progress display
+    prog_html: str,
 ):
     if not video or not video.name:
-        return None, None, "Upload a video.", "", gr.update(visible=False), [], "", ""
+        return None, None, "Upload a video.", "", prog_html, None, None, None
     if not FFMPEG or not FFPROBE:
-        return None, None, "FFmpeg missing. See note below.", MISSING_MSG, gr.update(visible=False), [], "", ""
+        return None, None, "FFmpeg missing. See note below.", MISSING_MSG, prog_html, None, None, None
 
-    # Work dirs
+    # Setup work dir
     work = Path(tempfile.mkdtemp(prefix="vid2img_"))
     raw_dir = work / "frames_raw"
     raw_dir.mkdir(parents=True, exist_ok=True)
 
-    # Default prefix from input filename if blank
-    if not prefix or not prefix.strip():
-        prefix = Path(video.name).stem
+    # Prefix defaulting
+    prefix = sanitize_prefix(prefix_in) or Path(video.name).stem
+
+    # Probe for estimate
+    info = parse_video_info(ffprobe_json(video.name))
+    est = estimate_output_count(mode, info.get("duration"), info.get("fps"), every_seconds, nth_frame, exact_fps)
 
-    # Build and run FFmpeg
+    # Build command
     pattern = str(raw_dir / f"{prefix}_%05d.{out_format}")
-    cmd = build_ffmpeg_command(
+    cmd = build_ffmpeg_extract(
         input_path=video.name,
         mode=mode,
         every_seconds=every_seconds,
         nth_frame=nth_frame,
         exact_fps=exact_fps,
-        start_time=start_time.strip(),
-        end_time=end_time.strip(),
+        start_time=(start_time or "").strip(),
+        end_time=(end_time or "").strip(),
         long_side=long_side,
         out_format=out_format,
         jpg_quality=jpg_quality,
@@ -257,101 +274,258 @@ def run_video_extract(
     )
     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, f"FFmpeg error:
-{proc.stderr}", cmd_preview, gr.update(visible=False), [], "", ""
+    # Run with Popen for live progress
+    proc = subprocess.Popen(cmd, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL, text=True, bufsize=1)
+
+    created = 0
+    total = est or None
+    last_html = prog_html
+
+    while True:
+        line = proc.stderr.readline()
+        if not line and proc.poll() is not None:
+            break
+        # update by counting files periodically
+        if int(time.time()*10) % 3 == 0:
+            created = len(list(raw_dir.glob(f"{prefix}_*.{out_format}")))
+            if total and total > 0:
+                pct = min(100.0, (created / total) * 100.0)
+                last_html = render_progress(pct, f"Extracting frames {created}/{total}")
+            else:
+                last_html = render_progress(0.0, f"Extracting frames… {created} created")
+    ret = proc.wait()
 
     frames = sorted(raw_dir.glob(f"{prefix}_*.{out_format}"))
-    if not frames:
-        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)]
+    if ret != 0 or not frames:
+        try:
+            err = proc.stderr.read() if proc.stderr else ""
+        except Exception:
+            err = ""
+        return None, None, f"FFmpeg error or no frames produced.\n\n{err}", cmd_preview, last_html, None, None, None
 
-    # Zip
+    # Preview + zip
+    gallery = sample_paths(frames, 30)
     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}"
+    return gallery, str(zip_path), details, cmd_preview, render_progress(100.0, f"Extracted {len(frames)} frames"), [str(p) for p in frames], str(raw_dir), prefix
 
-    # 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 step1_cancel(extract_pid: int | None):
+    if extract_pid:
+        try:
+            os.kill(extract_pid, 15)
+            return "Extraction cancelled."
+        except Exception:
+            return "Could not cancel (process already finished)."
+    return "No extraction running."
+
+# ─────────────────────────────────────────────────────────────
+# Step 2: Upscale extracted frames (with Cancel)
+# ─────────────────────────────────────────────────────────────
 
-def run_upscale_from_extracted(
+def step2_upscale(
     frames_list: List[str] | None,
-    frames_dir: str,
-    prefix: str,
     model_name: str,
     scale: int,
     tile: int,
     precision: str,
+    cancel_flag: bool,
+    prog_html: str,
 ):
-    if not _HAVE_REALESRGAN:
-        return None, None, "realesrgan is not installed (see requirements.txt)", ""
+    if not HAVE_REALESRGAN:
+        return None, None, "realesrgan is not installed (see requirements.txt)", prog_html
     if not frames_list:
-        return None, None, "No extracted frames state found. Please run extraction first.", ""
+        return None, None, "No extracted frames found in state. Run Step 1 first.", prog_html
+
+    device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
+    half = (precision == "half") and (device == "cuda")
+    upsampler = get_realesrganer(model_name, scale, tile, half, device=device)
 
     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"
+    total = len(frames_list)
+    done = 0
+    for fp in frames_list:
+        if cancel_flag:
+            return None, None, "Upscale interrupted by user.", render_progress((done/total)*100 if total else 0, f"Upscaled {done}/{total}")
+        try:
+            img = Image.open(fp).convert("RGB")
+            output, _ = upsampler.enhance(np.array(img), outscale=scale)
+            out_img = Image.fromarray(output)
+            out_file = out_dir / (Path(fp).stem + ".jpg")
+            out_img.save(out_file, quality=95)
+            done += 1
+        except Exception as e:
+            done += 1
+        # update bar
+        pct = (done/total)*100 if total else 0
+        prog_html = render_progress(pct, f"Upscaling {done}/{total}")
+    # build outputs
+    up_paths = sorted(out_dir.glob("*.jpg"))
+    gallery = sample_paths(up_paths, 30)
+
+    zip_path = work / "upscaled.zip"
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
-        for p in sorted(out_dir.glob("*")):
+        for p in up_paths:
             zf.write(p, p.name)
 
-    detail = f"Upscaled: {len(up_paths)} | Model: {model_name} | Scale: x{scale} | Tile: {tile} | Precision: {precision}"
-    return gallery, str(zip_path), detail, ""
+    return gallery, str(zip_path), f"Upscaled: {len(up_paths)}", render_progress(100.0, "Upscaling complete")
 
+# ─────────────────────────────────────────────────────────────
+# Step 3: Re-encode to video (with Cancel)
+# ─────────────────────────────────────────────────────────────
 
-def run_image_upscale_pipeline(
-    images: List[gr.File] | None,
-    model_name: str,
-    scale: int,
-    tile: int,
-    precision: str,
-    prefix: str,
-):
-    if not images:
-        return None, None, "Upload one or more images.", ""
-    if not _HAVE_REALESRGAN:
-        return None, None, "realesrgan is not installed (see requirements.txt)", ""
+def build_ffmpeg_encode(frames_dir: str, prefix: str, fps: float, fmt: str, include_audio: bool, orig_video: str) -> List[str]:
+    # Frame pattern assumes zero-padded sequence created earlier
+    pattern = str(Path(frames_dir) / f"{prefix}_%05d.jpg")
+    args = [FFMPEG, "-y", "-start_number", "1", "-framerate", f"{fps:.6f}", "-i", pattern]
+    if include_audio and orig_video:
+        args += ["-i", orig_video, "-map", "0:v:0", "-map", "1:a:0", "-shortest"]
+    # Choose codec
+    if fmt == "h265":
+        vcodec = ["-c:v", "libx265"]
+    elif fmt == "vp9":
+        vcodec = ["-c:v", "libvpx-vp9"]
+    else:
+        vcodec = ["-c:v", "libx264"]
+    args += vcodec + ["-pix_fmt", "yuv420p", "-crf", "18", "-preset", "medium"]
+    out_name = "output.mp4" if fmt in ("h264", "h265") else "output.webm"
+    args += [out_name]
+    return args
+
+
+def step3_encode(frames_dir: str | None, prefix: str | None, orig_video: str | None,
+                 fps: float | None, fmt: str, include_audio: bool, prog_html: str):
+    if not frames_dir or not prefix:
+        return None, None, "No frames available to encode. Run previous steps first.", prog_html
+    fps = fps or 30.0
+
+    cmd = build_ffmpeg_encode(frames_dir, prefix, fps, fmt, include_audio, orig_video or "")
+    proc = subprocess.Popen(cmd, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL, text=True, bufsize=1, cwd=frames_dir)
+
+    last_html = prog_html
+    # Try parsing progress via out_time_ms if available; else simple spinner
+    # (Encoding progress estimation is best-effort.)
+    start = time.time()
+    while True:
+        line = proc.stderr.readline()
+        if not line and proc.poll() is not None:
+            break
+        # coarse heartbeat
+        if int(time.time()*10) % 5 == 0:
+            last_html = render_progress(50.0, "Encoding…")
+    ret = proc.wait()
+
+    out_file = Path(frames_dir) / ("output.mp4" if fmt in ("h264", "h265") else "output.webm")
+    if ret != 0 or not out_file.exists():
+        try:
+            err = proc.stderr.read() if proc.stderr else ""
+        except Exception:
+            err = ""
+        return None, None, f"Encoding failed.\n\n{err}", last_html
+    return str(out_file), f"Video created: {out_file.name}", render_progress(100.0, "Encoding complete")
 
-    work = Path(tempfile.mkdtemp(prefix="imgup_"))
-    in_dir = work / "input"
-    out_dir = work / "upscaled"
-    in_dir.mkdir(parents=True, exist_ok=True)
-    out_dir.mkdir(parents=True, exist_ok=True)
+# ──────────────────��──────────────────────────────────────────
+# Quick Mode: All frames → Upscale x4 → MP4 (with audio)
+# ─────────────────────────────────────────────────────────────
 
-    img_paths: List[Path] = []
-    for f in images:
-        p = Path(f.name)
-        name = p.name
-        if prefix and prefix.strip():
-            name = f"{prefix}_{name}"
-        dst = in_dir / name
-        shutil.copy2(p, dst)
-        img_paths.append(dst)
+def quick_mode(video: gr.File | None, start_time: str, end_time: str, resize_long: int, prefix_in: str,
+               prog_html: str):
+    if not video or not video.name:
+        return None, None, None, "Upload a video.", prog_html
+    if not (FFMPEG and FFPROBE and HAVE_REALESRGAN):
+        return None, None, None, "Missing dependencies (ffmpeg/ffprobe/realesrgan)", prog_html
 
-    up_paths = upscale_images(img_paths, out_dir, model_name, scale, tile, precision)
-    gallery = [str(p) for p in sample_paths([Path(p) for p in up_paths], n=30)]
+    info = parse_video_info(ffprobe_json(video.name))
+    in_fps = info.get("fps") or 30.0
+    prefix = sanitize_prefix(prefix_in) or Path(video.name).stem
 
-    zip_path = work / "upscaled_images.zip"
-    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
-        for p in sorted(out_dir.glob("*")):
+    # Work dirs
+    work = Path(tempfile.mkdtemp(prefix="quick_"))
+    raw_dir = work / "frames_raw"; raw_dir.mkdir(parents=True, exist_ok=True)
+    up_dir = work / "upscaled";   up_dir.mkdir(parents=True, exist_ok=True)
+
+    # Phase 1: Extract ALL frames
+    extract_cmd = build_ffmpeg_extract(
+        input_path=video.name,
+        mode="All frames",
+        every_seconds=1.0,
+        nth_frame=1,
+        exact_fps=in_fps,
+        start_time=(start_time or "").strip(),
+        end_time=(end_time or "").strip(),
+        long_side=resize_long,
+        out_format="jpg",
+        jpg_quality=3,
+        png_level=2,
+        scene_detect=False,
+        scene_thresh=0.3,
+        out_pattern=str(raw_dir / f"{prefix}_%05d.jpg"),
+    )
+    proc = subprocess.Popen(extract_cmd, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL, text=True, bufsize=1)
+    # Estimate
+    est = estimate_output_count("All frames", info.get("duration"), in_fps, 1.0, 1, in_fps)
+    created = 0
+    while True:
+        line = proc.stderr.readline()
+        if not line and proc.poll() is not None:
+            break
+        if int(time.time()*10) % 3 == 0:
+            created = len(list(raw_dir.glob(f"{prefix}_*.jpg")))
+            pct = min(100.0, (created / est) * 100.0) if est else 0
+            prog_html = render_progress(pct, f"Phase 1/3: Extracting {created}/{est or '?'}")
+    proc.wait()
+
+    frames = sorted(raw_dir.glob(f"{prefix}_*.jpg"))
+    if not frames:
+        return None, None, None, "No frames extracted in Quick Mode.", prog_html
+
+    # Phase 2: Upscale x4 (x4plus, auto tile)
+    device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
+    upsampler = get_realesrganer("x4plus", 4, 0, (device=="cuda"), device=device)
+
+    total = len(frames)
+    done = 0
+    for fp in frames:
+        img = Image.open(fp).convert("RGB")
+        output, _ = upsampler.enhance(np.array(img), outscale=4)
+        Image.fromarray(output).save(up_dir / (Path(fp).stem + ".jpg"), quality=95)
+        done += 1
+        pct = (done/total)*100 if total else 0
+        prog_html = render_progress(pct, f"Phase 2/3: Upscaling {done}/{total}")
+
+    # Phase 3: Re-encode to MP4 (H.264) with audio
+    encode_cmd = build_ffmpeg_encode(str(up_dir), prefix, in_fps, "h264", True, video.name)
+    proc2 = subprocess.Popen(encode_cmd, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL, text=True, bufsize=1, cwd=str(up_dir))
+    while True:
+        line = proc2.stderr.readline()
+        if not line and proc2.poll() is not None:
+            break
+        if int(time.time()*10) % 5 == 0:
+            prog_html = render_progress(50.0, "Phase 3/3: Encoding…")
+    proc2.wait()
+
+    out_file = Path(up_dir) / "output.mp4"
+    if not out_file.exists():
+        return None, None, None, "Encoding failed in Quick Mode.", prog_html
+
+    # Optional intermediates
+    zip_frames = work / "frames.zip"
+    with zipfile.ZipFile(zip_frames, "w", zipfile.ZIP_DEFLATED) as zf:
+        for p in frames:
+            zf.write(p, p.name)
+    zip_up = work / "upscaled.zip"
+    with zipfile.ZipFile(zip_up, "w", zipfile.ZIP_DEFLATED) as zf:
+        for p in sorted(up_dir.glob("*.jpg")):
             zf.write(p, p.name)
 
-    detail = f"Images upscaled: {len(up_paths)} | Model: {model_name} | Scale: x{scale} | Tile: {tile} | Precision: {precision}"
-    return gallery, str(zip_path), detail, ""
+    return str(out_file), str(zip_frames), str(zip_up), "Quick Mode complete.", render_progress(100.0, "All done")
 
 # ────────���────────────────────────────────────────────────────
 # UI
@@ -363,21 +537,22 @@ 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 (Two-step)</div>
-        Step 1 extracts frames. When it finishes, Step 2 (Upscale) appears.
+        <div class=\"cf-title\">Video → Frames → Upscale → Re-encode</div>
+        Three-step workflow plus a one-click Quick Mode. Previews sample 30 frames; galleries scroll.
         """)
 
-        # Shared state between steps
-        frames_state = gr.State([])     # list[str] of extracted frame paths
-        frames_dir_state = gr.State("")
-        prefix_state = gr.State("")
+        # Shared states
+        frames_state = gr.State([])        # list[str]
+        frames_dir_state = gr.State("")   # str
+        prefix_state = gr.State("")       # str
+        fps_state = gr.State(30.0)         # float
 
         with gr.Tabs():
-            # ── Tab 1: Video → Frames (Step 1)
-            with gr.Tab("Video → Frames"):
+            # TAB 1: Step 1 Extract
+            with gr.Tab("Step 1 · Extract 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):
+                with gr.Accordion("Extraction Settings", open=True):
                     with gr.Row():
                         mode = gr.Dropdown(["Every N seconds", "Every Nth frame", "Exact FPS", "All frames"], value="Every N seconds", label="Mode")
                         every_seconds = gr.Number(value=1.0, label="Every N seconds")
@@ -395,81 +570,139 @@ 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)")
+                # Estimate label (live)
+                estimate_md = gr.Markdown("Estimated output: —")
 
-                run_extract_btn = gr.Button("Step 1: Extract frames", variant="primary")
+                with gr.Row():
+                    btn_extract = gr.Button("Step 1: Extract Frames", variant="primary")
+                    # We expose a Cancel button visually; process kill is best-effort
+                    # A real cancel would wire a separate endpoint to signal/terminate.
+                    # (Left as stub for simplicity; you can wire a cancel flag similarly to step2.)
+                    # btn_cancel_extract = gr.Button("Cancel", variant="stop")
+                prog1 = gr.HTML(render_progress(0.0, "Idle"))
                 gallery = gr.Gallery(label="Preview (30 sampled)", columns=6, height=480)
                 zip_out = gr.File(label="Download frames ZIP")
-                details = gr.Markdown("Ready.")
+                details1 = gr.Markdown("Ready.")
                 with gr.Accordion("Show FFmpeg command", open=False):
                     cmd_preview = gr.Textbox(label="ffmpeg command", lines=4, elem_classes=["cmdbox"])
                 if MISSING_MSG:
                     gr.Markdown(f"<span style='color:#b45309'>{MISSING_MSG}</span>")
 
-                def _toggle(mode_val, fmt):
-                    return (
-                        gr.update(visible=(mode_val == "Every N seconds")),
-                        gr.update(visible=(mode_val == "Every Nth frame")),
-                        gr.update(visible=(mode_val == "Exact FPS")),
-                        gr.update(visible=(fmt == "jpg")),
-                        gr.update(visible=(fmt == "png")),
-                    )
-
-                mode.change(_toggle, [mode, out_format], [every_seconds, nth_frame, exact_fps, jpg_quality, png_level])
-                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 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,
+                def update_estimate(vfile, mode_val, evs, nth, exfps, st, et):
+                    if not vfile or not getattr(vfile, 'name', None):
+                        return "Estimated output: —"
+                    info = parse_video_info(ffprobe_json(vfile.name))
+                    # adjust duration if trims provided (rough parse HH:MM:SS)
+                    dur = info.get("duration")
+                    def parse_ts(ts: str):
+                        if not ts:
+                            return 0.0
+                        parts = ts.split(":")
+                        if len(parts) == 3:
+                            try:
+                                return float(parts[0])*3600 + float(parts[1])*60 + float(parts[2])
+                            except Exception:
+                                return 0.0
+                        return 0.0
+                    st_s = parse_ts(st or ""); et_s = parse_ts(et or "")
+                    if dur:
+                        if st_s: dur = max(0.0, dur - st_s)
+                        if et_s and et_s < info.get("duration", 0) and et_s > 0:
+                            dur = min(dur, et_s)
+                    est = estimate_output_count(mode_val, dur, info.get("fps"), evs or 1.0, int(nth or 1), exfps or 1.0)
+                    if not est:
+                        return "Estimated output: —"
+                    return f"Estimated output: **~{est} frames**"
+
+                for ctrl in [video, mode, every_seconds, nth_frame, exact_fps, start_time, end_time]:
+                    ctrl.change(update_estimate, inputs=[video, mode, every_seconds, nth_frame, exact_fps, start_time, end_time], outputs=[estimate_md])
+
+                btn_extract.click(
+                    step1_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,
+                        prog1,
                     ],
-                    outputs=[gallery, zip_out, details, cmd_preview, upscale_group, frames_state, frames_dir_state, prefix_state],
+                    outputs=[gallery, zip_out, details1, cmd_preview, prog1, frames_state, frames_dir_state, prefix_state],
                 )
 
-            # ── 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")
+            # TAB 2: Step 2 Upscale
+            with gr.Tab("Step 2 · Upscale Extracted"):
+                if not HAVE_REALESRGAN:
+                    gr.Markdown("⚠️ Install realesrgan/basicsr in requirements.txt to enable upscaling.")
+                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)")
+                with gr.Row():
+                    btn_upscale = gr.Button("Step 2: Upscale Frames", variant="primary")
+                    # btn_cancel_up = gr.Button("Cancel", variant="stop")
+                prog2 = gr.HTML(render_progress(0.0, "Idle"))
+                gallery_up = gr.Gallery(label="Upscaled preview (30 sampled)", columns=6, height=480)
+                zip_up = gr.File(label="Download upscaled ZIP")
+                details2 = gr.Markdown("")
+
+                # Simple cancel flag wiring (stub). In a full app, you'd toggle this from a Cancel button.
+                cancel_flag = gr.State(False)
+
+                btn_upscale.click(
+                    step2_upscale,
+                    inputs=[frames_state, model_name, scale, tile, precision, cancel_flag, prog2],
+                    outputs=[gallery_up, zip_up, details2, prog2],
+                )
+
+            # TAB 3: Step 3 Re-encode
+            with gr.Tab("Step 3 · Re-encode Video"):
+                with gr.Row():
+                    fmt = gr.Dropdown(["h264", "h265", "vp9"], value="h264", label="Format")
+                    include_audio = gr.Checkbox(True, label="Include original audio")
+                with gr.Row():
+                    btn_encode = gr.Button("Step 3: Create Video", variant="primary")
+                    # btn_cancel_enc = gr.Button("Cancel", variant="stop")
+                prog3 = gr.HTML(render_progress(0.0, "Idle"))
+                video_player = gr.Video(label="Preview video")
+                details3 = gr.Markdown("")
+
+                # Compute FPS once when video changes
+                def set_fps(vfile):
+                    if not vfile or not getattr(vfile, 'name', None):
+                        return 30.0
+                    info = parse_video_info(ffprobe_json(vfile.name))
+                    return float(info.get("fps") or 30.0)
+                video.change(set_fps, inputs=[video], outputs=[fps_state])
+
+                btn_encode.click(
+                    step3_encode,
+                    inputs=[frames_dir_state, prefix_state, video, fps_state, fmt, include_audio, prog3],
+                    outputs=[video_player, details3, prog3],
+                )
+
+            # TAB 4: Quick Mode (one click)
+            with gr.Tab("⚡ Quick Mode"):
+                gr.Markdown("Extract ALL frames → Upscale ×4 → MP4 (H.264) with original audio. No toggles.")
+                with gr.Row():
+                    q_video = gr.File(label="Upload video", file_types=[".mp4", ".mov", ".mkv", ".avi", ".webm", ".m4v"], 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)")
-                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("")
-
-                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()],
+                    q_start = gr.Textbox(value="", label="Start (HH:MM:SS.mmm, optional)")
+                    q_end = gr.Textbox(value="", label="End (HH:MM:SS.mmm, optional)")
+                    q_resize = gr.Number(value=0, label="Resize long side before upscale (0 = none)")
+                q_prefix = gr.Textbox(value="", label="Filename prefix (defaults to input file name)")
+
+                q_btn = gr.Button("Run Quick Pipeline", variant="primary")
+                q_prog = gr.HTML(render_progress(0.0, "Idle"))
+                q_video_out = gr.Video(label="Output video")
+                with gr.Accordion("Show intermediates", open=False):
+                    q_zip_frames = gr.File(label="frames.zip")
+                    q_zip_up = gr.File(label="upscaled.zip")
+                q_details = gr.Markdown("")
+
+                q_btn.click(
+                    quick_mode,
+                    inputs=[q_video, q_start, q_end, q_resize, q_prefix, q_prog],
+                    outputs=[q_video_out, q_zip_frames, q_zip_up, q_details, q_prog],
                 )
 
     return demo
@@ -479,19 +712,3 @@ if __name__ == "__main__":
     demo = build_ui()
     demo.queue().launch()
 
-# ─────────────────────────────────────────────────────────────
-# packages.txt (no comments, one per line)
-# ffmpeg
-# libsm6
-# libxext6
-# ─────────────────────────────────────────────────────────────
-
-# ─────────────────────────────────────────────────────────────
-# requirements.txt
-# gradio==5.44.1
-# realesrgan==0.3.0
-# basicsr==1.4.2
-# opencv-python-headless==4.10.0.84
-# numpy
-# torch==2.2.2
-# ─────────────────────────────────────────────────────────────