Update app.py

app.py

@@ -1,13 +1,12 @@
 # =============================
 # app.py
 # 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
+# - Step 1: Extract frames (with live estimate & progress)
+# - Step 2: Upscale frames (now supports uploading your own images directly)
+# - Step 3: Re-encode frames (now supports uploading your own frames/ZIP and optional audio source)
+# - Quick Mode: One-click (All Frames → Upscale x4 → MP4 H.264 with audio)
+# - Previews: 30 frames sampled evenly; scrollable galleries
 # - Prefix defaults to input video filename if left blank
-# - Direct image-upscale tab also included
 # =============================
 
 import os
@@ -39,17 +38,23 @@ FFPROBE = _which("ffprobe")
 
 if not FFMPEG or not FFPROBE:
     MISSING_MSG = (
-        "⚠️ FFmpeg not found. Add a 'packages.txt' with exactly:\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 = ""
 
+# Try to import Real-ESRGAN stack
 try:
     from realesrgan import RealESRGANer
     from basicsr.archs.rrdbnet_arch import RRDBNet
     HAVE_REALESRGAN = True
-except Exception:
+except Exception as e:
     HAVE_REALESRGAN = False
+    REAL_ERR = str(e)
 
 # ─────────────────────────────────────────────────────────────
 # Helpers
@@ -182,7 +187,7 @@ def build_ffmpeg_extract(
 
 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. See requirements.txt")
+        raise RuntimeError("realesrgan is not installed. See requirements.txt (realesrgan, basicsr, torch, numpy, scipy, scikit-image).")
     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
@@ -205,19 +210,12 @@ def get_realesrganer(model_name: str, scale: int, tile: int, half: bool, device:
     )
     return upsampler
 
-# ─────────────────────────────────────────────────────────────
-# Progress UI helper (HTML bar)
-# ─────────────────────────────────────────────────────────────
 
 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>'''
+    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>'''
 
-# ─────────────────────────────────────────────────────────────
-# Step 1: Extract frames (with Cancel)
-# ─────────────────────────────────────────────────────────────
+# ───────────────── Extraction (Step 1)
 
 def step1_extract(
     video: gr.File | None,
@@ -234,7 +232,6 @@ def step1_extract(
     scene_detect: bool,
     scene_thresh: float,
     prefix_in: str,
-    # progress display
     prog_html: str,
 ):
     if not video or not video.name:
@@ -242,19 +239,15 @@ def step1_extract(
     if not FFMPEG or not FFPROBE:
         return None, None, "FFmpeg missing. See note below.", MISSING_MSG, prog_html, None, None, None
 
-    # Setup work dir
     work = Path(tempfile.mkdtemp(prefix="vid2img_"))
     raw_dir = work / "frames_raw"
     raw_dir.mkdir(parents=True, exist_ok=True)
 
-    # 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 command
     pattern = str(raw_dir / f"{prefix}_%05d.{out_format}")
     cmd = build_ffmpeg_extract(
         input_path=video.name,
@@ -274,25 +267,22 @@ def step1_extract(
     )
     cmd_preview = " ".join([s if " " not in s else f'"{s}"' for s in cmd])
 
-    # 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}")
+                last_html = render_progress(pct, f"Extracting {created}/{total}")
             else:
-                last_html = render_progress(0.0, f"Extracting frames… {created} created")
+                last_html = render_progress(0.0, f"Extracting… {created} created")
     ret = proc.wait()
 
     frames = sorted(raw_dir.glob(f"{prefix}_*.{out_format}"))
@@ -301,9 +291,10 @@ def step1_extract(
             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
+        return None, None, f"FFmpeg error or no frames produced.
+
+{err}", cmd_preview, last_html, None, None, None
 
-    # Preview + zip
     gallery = sample_paths(frames, 30)
     zip_path = work / "frames.zip"
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
@@ -313,19 +304,22 @@ def step1_extract(
     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
 
+# ───────────────── Upscale (Step 2) — supports uploaded images OR frames from Step 1
 
-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."
+def save_uploaded_images(files: List[gr.File] | None, prefix: str = "upload") -> Tuple[List[Path], Path]:
+    tmp = Path(tempfile.mkdtemp(prefix="imgup_"))
+    in_dir = tmp / "input"; in_dir.mkdir(parents=True, exist_ok=True)
+    paths: List[Path] = []
+    if not files:
+        return paths, in_dir
+    for f in files:
+        src = Path(f.name)
+        name = f"{prefix}_{src.name}"
+        dst = in_dir / name
+        shutil.copy2(src, dst)
+        paths.append(dst)
+    return paths, in_dir
 
-# ─────────────────────────────────────────────────────────────
-# Step 2: Upscale extracted frames (with Cancel)
-# ─────────────────────────────────────────────────────────────
 
 def step2_upscale(
     frames_list: List[str] | None,
@@ -333,43 +327,51 @@ def step2_upscale(
     scale: int,
     tile: int,
     precision: str,
-    cancel_flag: bool,
     prog_html: str,
+    uploaded_imgs: List[gr.File] | None,
 ):
     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 found in state. Run Step 1 first.", prog_html
+        msg = "Real-ESRGAN not available. Ensure requirements.txt includes: --prefer-binary, numpy==1.26.4, scipy==1.11.4, scikit-image==0.22.0, opencv-python-headless, torch==2.2.2, realesrgan==0.3.0, basicsr==1.4.2, pillow, gradio."
+        return None, None, msg, prog_html
+
+    # decide source: uploaded images take priority, else frames from step 1
+    if uploaded_imgs and len(uploaded_imgs) > 0:
+        img_paths, _ = save_uploaded_images(uploaded_imgs, prefix="up")
+        src_paths = [str(p) for p in img_paths]
+    else:
+        src_paths = frames_list or []
+
+    if not src_paths:
+        return None, None, "No images provided. Upload files or 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)
+    out_dir = work / "upscaled"; out_dir.mkdir(parents=True, exist_ok=True)
 
-    total = len(frames_list)
+    total = len(src_paths)
     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}")
+    up_paths: List[Path] = []
+    for fp in src_paths:
         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
+            up_paths.append(out_file)
+        except Exception:
+            pass
+        done += 1
         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)
 
+    if not up_paths:
+        return None, None, "Upscaling produced no outputs.", prog_html
+
+    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 up_paths:
@@ -377,17 +379,44 @@ def step2_upscale(
 
     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 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")
+# ───────────────── Encode (Step 3) — supports uploaded frames/ZIP & optional audio source
+
+def prepare_frames_from_upload(files: List[gr.File] | None, prefix: str = "enc") -> Tuple[Optional[str], Optional[str]]:
+    if not files:
+        return None, None
+    work = Path(tempfile.mkdtemp(prefix="enc_"))
+    frames_dir = work / "frames"; frames_dir.mkdir(parents=True, exist_ok=True)
+    detected_prefix = None
+
+    # If a single ZIP is uploaded, unzip
+    if len(files) == 1 and Path(files[0].name).suffix.lower() == ".zip":
+        with zipfile.ZipFile(files[0].name, "r") as zf:
+            zf.extractall(frames_dir)
+        # try detect a prefix
+        imgs = sorted(frames_dir.glob("*.jpg")) + sorted(frames_dir.glob("*.png"))
+        if imgs:
+            detected_prefix = Path(imgs[0]).stem.split("_")[0]
+        return str(frames_dir), detected_prefix or prefix
+
+    # else, copy images directly
+    counter = 1
+    for f in files:
+        src = Path(f.name)
+        if src.suffix.lower() not in [".jpg", ".jpeg", ".png"]:
+            continue
+        dst = frames_dir / f"{prefix}_{counter:05d}{src.suffix.lower()}"
+        shutil.copy2(src, dst)
+        counter += 1
+    return str(frames_dir), prefix
+
+
+def build_ffmpeg_encode(frames_dir: str, prefix: str, fps: float, fmt: str, include_audio: bool, orig_video: str | None) -> List[str]:
+    pattern_jpg = Path(frames_dir) / f"{prefix}_%05d.jpg"
+    pattern_png = Path(frames_dir) / f"{prefix}_%05d.png"
+    pattern = str(pattern_jpg if pattern_jpg.exists() else pattern_png)
     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":
@@ -400,24 +429,30 @@ def build_ffmpeg_encode(frames_dir: str, prefix: str, fps: float, fmt: str, incl
     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):
+def step3_encode(frames_dir_state: str | None, prefix_state: str | None, orig_video: gr.File | None,
+                 fps: float | None, fmt: str, include_audio: bool, prog_html: str,
+                 uploaded_frames: List[gr.File] | None, uploaded_audio_video: gr.File | None):
+    # Choose frames source: uploaded takes priority
+    frames_dir = frames_dir_state
+    prefix = prefix_state
+    if uploaded_frames and len(uploaded_frames) > 0:
+        frames_dir, detected = prepare_frames_from_upload(uploaded_frames, prefix or "enc")
+        if detected:
+            prefix = detected
     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
+        return None, None, "No frames available. Upload frames (ZIP/images) or run Step 1.", prog_html
+
+    fps = float(fps or 30.0)
+    orig_path = uploaded_audio_video.name if uploaded_audio_video else (orig_video.name if orig_video else None)
 
-    cmd = build_ffmpeg_encode(frames_dir, prefix, fps, fmt, include_audio, orig_video or "")
+    cmd = build_ffmpeg_encode(frames_dir, prefix, fps, fmt, include_audio, orig_path)
     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()
@@ -428,30 +463,28 @@ def step3_encode(frames_dir: str | None, prefix: str | None, orig_video: str | N
             err = proc.stderr.read() if proc.stderr else ""
         except Exception:
             err = ""
-        return None, None, f"Encoding failed.\n\n{err}", last_html
+        return None, None, f"Encoding failed.
+
+{err}", last_html
     return str(out_file), f"Video created: {out_file.name}", render_progress(100.0, "Encoding complete")
 
-# ─────────────────────────────────────────────────────────────
-# Quick Mode: All frames → Upscale x4 → MP4 (with audio)
-# ─────────────────────────────────────────────────────────────
+# ───────────────── Quick Mode — one click: All frames → Upscale ×4 → MP4 (audio)
 
-def quick_mode(video: gr.File | None, start_time: str, end_time: str, resize_long: int, prefix_in: str,
-               prog_html: str):
+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
+        return None, None, None, "Missing deps (ffmpeg/ffprobe/realesrgan). See requirements.txt.", prog_html
 
     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
 
-    # 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 all frames
     extract_cmd = build_ffmpeg_extract(
         input_path=video.name,
         mode="All frames",
@@ -469,7 +502,6 @@ def quick_mode(video: gr.File | None, start_time: str, end_time: str, resize_lon
         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:
@@ -486,7 +518,7 @@ def quick_mode(video: gr.File | None, start_time: str, end_time: str, resize_lon
     if not frames:
         return None, None, None, "No frames extracted in Quick Mode.", prog_html
 
-    # Phase 2: Upscale x4 (x4plus, auto tile)
+    # Upscale x4
     device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
     upsampler = get_realesrganer("x4plus", 4, 0, (device=="cuda"), device=device)
 
@@ -500,7 +532,7 @@ def quick_mode(video: gr.File | None, start_time: str, end_time: str, resize_lon
         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 MP4 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:
@@ -515,7 +547,7 @@ def quick_mode(video: gr.File | None, start_time: str, end_time: str, resize_lon
     if not out_file.exists():
         return None, None, None, "Encoding failed in Quick Mode.", prog_html
 
-    # Optional intermediates
+    # Intermediates
     zip_frames = work / "frames.zip"
     with zipfile.ZipFile(zip_frames, "w", zipfile.ZIP_DEFLATED) as zf:
         for p in frames:
@@ -527,9 +559,7 @@ def quick_mode(video: gr.File | None, start_time: str, end_time: str, resize_lon
 
     return str(out_file), str(zip_frames), str(zip_up), "Quick Mode complete.", render_progress(100.0, "All done")
 
-# ─────────────────────────────────────────────────────────────
-# UI
-# ─────────────────────────────────────────────────────────────
+# ───────────────── UI
 
 def build_ui():
     with gr.Blocks(theme=gr.themes.Soft(), css="""
@@ -538,17 +568,17 @@ def build_ui():
     """) as demo:
         gr.Markdown("""
         <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.
+        Three-step workflow + Quick Mode. Step 2/3 now accept your own uploaded files as inputs.
         """)
 
-        # Shared states
+        # Shared states (from Step 1)
         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: Step 1 Extract
+            # STEP 1
             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")
@@ -570,15 +600,9 @@ 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: —")
-
                 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")
@@ -592,17 +616,13 @@ def build_ui():
                     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
+                        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
+                            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:
@@ -610,9 +630,7 @@ def build_ui():
                         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**"
+                    return f"Estimated output: **~{est} frames**" if est else "Estimated output: —"
 
                 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])
@@ -628,59 +646,59 @@ def build_ui():
                     outputs=[gallery, zip_out, details1, cmd_preview, prog1, frames_state, frames_dir_state, prefix_state],
                 )
 
-            # TAB 2: Step 2 Upscale
-            with gr.Tab("Step 2 · Upscale Extracted"):
+            # STEP 2 — Upscale
+            with gr.Tab("Step 2 · Upscale Frames"):
                 if not HAVE_REALESRGAN:
-                    gr.Markdown("⚠️ Install realesrgan/basicsr in requirements.txt to enable upscaling.")
+                    gr.Markdown("⚠️ Upscaling disabled. Install dependencies in requirements.txt (see notes in code). Error: " + (REAL_ERR if 'REAL_ERR' in globals() else ""))
+                gr.Markdown("Use frames from Step 1 **or** upload images below.")
+                imgs_override = gr.Files(label="Upload images to upscale (JPG/PNG)", file_types=[".jpg", ".jpeg", ".png"], type="filepath")
                 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")
+                    btn_upscale = gr.Button("Step 2: Upscale", variant="primary")
                 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],
+                    inputs=[frames_state, model_name, scale, tile, precision, prog2, imgs_override],
                     outputs=[gallery_up, zip_up, details2, prog2],
                 )
 
-            # TAB 3: Step 3 Re-encode
+            # STEP 3 — Re-encode
             with gr.Tab("Step 3 · Re-encode Video"):
+                gr.Markdown("Use frames from Step 1 **or** upload a frames ZIP / images. Optionally provide a video for audio track.")
+                uploaded_frames = gr.Files(label="Upload frames (ZIP or images)", type="filepath")
+                uploaded_audio = gr.File(label="Optional: video/audio source for audio track", file_types=[".mp4", ".mov", ".mkv", ".webm", ".mp3", ".wav"], type="filepath")
                 with gr.Row():
                     fmt = gr.Dropdown(["h264", "h265", "vp9"], value="h264", label="Format")
-                    include_audio = gr.Checkbox(True, label="Include original audio")
+                    include_audio = gr.Checkbox(True, label="Include audio if available")
                 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)
+                # capture FPS from the original step1 video when it changes
                 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],
+                    inputs=[frames_dir_state, prefix_state, video, fps_state, fmt, include_audio, prog3, uploaded_frames, uploaded_audio],
                     outputs=[video_player, details3, prog3],
                 )
 
-            # TAB 4: Quick Mode (one click)
+            # QUICK MODE
             with gr.Tab("⚡ Quick Mode"):
                 gr.Markdown("Extract ALL frames → Upscale ×4 → MP4 (H.264) with original audio. No toggles.")
                 with gr.Row():
@@ -710,5 +728,4 @@ def build_ui():
 
 if __name__ == "__main__":
     demo = build_ui()
-    demo.queue().launch()
-
+    demo.queue().launch()