Update app.py

app.py

@@ -518,9 +518,11 @@ def step1_extract(
     prog_html: str,
 ):
     if not video or not video.name:
-        return None, None, "Upload a video.", "", prog_html, None, None, None
+        yield None, None, "Upload a video.", "", prog_html, None, None, None
+        return
     if not FFMPEG or not FFPROBE:
-        return None, None, "FFmpeg missing. See note below.", MISSING_MSG, prog_html, None, None, None
+        yield None, None, "FFmpeg missing. See note below.", MISSING_MSG, prog_html, None, None, None
+        return
 
     work = Path(tempfile.mkdtemp(prefix="vid2img_"))
     raw_dir = work / "frames_raw"
@@ -528,9 +530,24 @@ def step1_extract(
 
     prefix = sanitize_prefix(prefix_in) or Path(video.name).stem
 
-    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)
+    # Duration for progress %
+    vinfo = parse_video_info(ffprobe_json(video.name))
+    full_duration = float(vinfo.get("duration") or 0.0)
+
+    # If trimming, adjust expected duration window
+    def _parse_ts(ts: str) -> float:
+        if not ts: return 0.0
+        h, m, s = ts.split(":") if ":" in ts else ("0","0",ts)
+        return float(h)*3600 + float(m)*60 + float(s)
+
+    st_s = _parse_ts((start_time or "").strip())
+    et_s = _parse_ts((end_time or "").strip())
+    if full_duration and st_s > 0:
+        full_duration = max(0.0, full_duration - st_s)
+    if full_duration and et_s > 0 and et_s < (vinfo.get("duration") or 0):
+        full_duration = max(0.0, min(full_duration, et_s))
 
+    # Build command
     pattern = str(raw_dir / f"{prefix}_%05d.{out_format}")
     cmd = build_ffmpeg_extract(
         input_path=video.name,
@@ -548,41 +565,67 @@ def step1_extract(
         scene_thresh=scene_thresh,
         out_pattern=pattern,
     )
+
+    # Inject progress reporting
+    # ffmpeg will write key=value lines including out_time to stderr
+    cmd = [cmd[0], "-progress", "pipe:2"] + cmd[1:]
     cmd_preview = " ".join([s if " " not in s else f'"{s}"' for s in cmd])
 
-    proc = subprocess.Popen(cmd, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL, text=True, bufsize=1)
+    proc = subprocess.Popen(
+        cmd, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL, text=True, bufsize=1
+    )
 
-    created = 0
-    total = est or None
-    last_html = prog_html
+    # Stream updates
+    last_pct = 0.0
+    gallery_preview = []
     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}_*.{out_format}")))
-            if total and total > 0:
-                pct = min(100.0, (created / total) * 100.0)
-                last_html = render_progress(pct, f"Extracting {created}/{total}")
-            else:
-                last_html = render_progress(0.0, f"Extracting… {created} created")
+
+        line = (line or "").strip()
+        # out_time is in HH:MM:SS.microsec
+        if line.startswith("out_time=") and full_duration > 0:
+            t = line.split("=", 1)[1]
+            # Convert HH:MM:SS.xx to seconds
+            try:
+                h, m, s = t.split(":")
+                secs = float(h) * 3600 + float(m) * 60 + float(s)
+            except Exception:
+                secs = 0.0
+            pct = max(0.0, min(100.0, (secs / full_duration) * 100.0))
+            if pct - last_pct >= 1.0 or pct in (0.0, 100.0):
+                last_pct = pct
+                # Lightweight live gallery sample for feel; not required
+                gallery_preview = sample_paths(sorted(raw_dir.glob(f"{prefix}_*.{out_format}"), key=_natural_key), 36)
+                yield gallery_preview, None, "Extracting…", cmd_preview, render_progress(pct, f"Extracting {pct:.0f}%"), None, str(raw_dir), prefix
+
     ret = proc.wait()
 
     frames = sorted(raw_dir.glob(f"{prefix}_*.{out_format}"), key=_natural_key)
 
-    # Adaptive preview
+    # Show all if ≤100, else sample 100
     if len(frames) <= 100:
-        gallery = [str(p) for p in frames]   # show all if small set
+        gallery = [str(p) for p in frames]
     else:
-        gallery = sample_paths(frames, 100)  # evenly sample 100
+        gallery = sample_paths(frames, 100)
 
     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)
 
+    if ret != 0 or not frames:
+        err = ""
+        try:
+            err = proc.stderr.read() if proc.stderr else ""
+        except Exception:
+            pass
+        yield gallery, None, f"Extraction failed.\n\n{err}", cmd_preview, render_progress(0.0, "Failed"), None, str(raw_dir), prefix
+        return
+
     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
+    yield 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
 
@@ -628,8 +671,10 @@ def step2_process_next_batch(
     up_src_paths, up_out_dir, up_done_idx, up_total,
     ui_model_name, outscale, tile, precision, denoise_strength, face_enhance, batch_size,
 ):
+    # Turn this into a generator that streams progress
     if not up_src_paths or not up_out_dir:
-        return None, None, "Load sources first.", render_progress(0.0, "Idle"), up_done_idx, up_out_dir
+        yield None, None, "Load sources first.", render_progress(0.0, "Idle"), up_done_idx, up_out_dir
+        return
 
     model_id = map_ui_model_to_internal(ui_model_name)
     scale = clamp_scale_for_model(int(outscale or 4), model_id)
@@ -667,10 +712,16 @@ def step2_process_next_batch(
         zip_file = _save_zip_of_dir(out_dir, zip_path)
         prog = render_progress(100.0, "All images processed")
         details = f"Done. Total upscaled: {len(list(out_dir.glob('*.jpg')))+len(list(out_dir.glob('*.png')))}"
-        return gallery, zip_file, details, prog, start, up_out_dir
+        yield gallery, zip_file, details, prog, start, up_out_dir
+        return
 
+    batch_paths = up_src_paths[start:end]
+    total_in_batch = len(batch_paths)
     processed_now = 0
-    for fp in up_src_paths[start:end]:
+
+    # For ETA
+    t0 = time.time()
+    for idx, fp in enumerate(batch_paths, start=1):
         try:
             with Image.open(fp) as im:
                 img = im.convert("RGB")
@@ -681,27 +732,41 @@ def step2_process_next_batch(
                         cv_img, has_aligned=False, only_center_face=False, paste_back=True
                     )
                 else:
+                    # denoise_strength only applies to general-x4v3, but harmless otherwise
                     output, _ = upsampler.enhance(cv_img, outscale=scale, denoise_strength=float(denoise_strength or 0.5))
 
                 Image.fromarray(output).save(out_dir / (Path(fp).stem + ".jpg"), quality=95)
 
         except Exception as e:
             print("Upscale error:", e)
-        processed_now += 1
 
-    next_idx = end
-    pct = int(round((next_idx / up_total) * 100)) if up_total else 0
-    label = (f"Processed {processed_now} image(s) this batch. "
-             f"{next_idx}/{up_total} done (x{scale}, model={ui_model_name}, "
-             f"denoise={denoise_strength}, face={face_enhance}).")
-    prog = render_progress(pct, f"Upscaling… {pct}%")
+        processed_now = idx
+        # Progress & ETA for THIS batch
+        pct_batch = (processed_now / total_in_batch) * 100.0
+        elapsed = time.time() - t0
+        secs_per_img = elapsed / max(1, processed_now)
+        remaining_imgs = total_in_batch - processed_now
+        eta = remaining_imgs * secs_per_img
+        label = (f"Batch: {processed_now}/{total_in_batch} · "
+                 f"~{eta:.1f}s ETA · global {start+processed_now}/{up_total} "
+                 f"(x{scale}, model={ui_model_name}, denoise={denoise_strength}, face={face_enhance})")
+
+        gallery = _build_gallery_from_dir(out_dir, 30)
+        zip_path = Path(out_dir.parent) / "upscaled.zip"
+        zip_file = _save_zip_of_dir(out_dir, zip_path)
+        yield gallery, zip_file, label, render_progress(pct_batch, f"Upscaling… {pct_batch:.0f}% (this batch)"), start+processed_now, up_out_dir
 
+    # Batch complete — final emit for this click
+    next_idx = end
     gallery = _build_gallery_from_dir(out_dir, 30)
     zip_path = Path(out_dir.parent) / "upscaled.zip"
     zip_file = _save_zip_of_dir(out_dir, zip_path)
-    return gallery, zip_file, label, prog, next_idx, up_out_dir
-
 
+    # Total (global) percentage across all sources
+    pct_global = (next_idx / up_total) * 100.0 if up_total else 100.0
+    final_label = (f"Processed batch {total_in_batch} image(s). "
+                   f"{next_idx}/{up_total} done (global {pct_global:.0f}%).")
+    yield gallery, zip_file, final_label, render_progress(pct_global, "Upscaling… (global)"), next_idx, up_out_dir
 
 
 def save_uploaded_images(files: List[gr.File] | None, prefix: str = "upload") -> Tuple[List[Path], Path]:
@@ -876,9 +941,17 @@ def build_ffmpeg_encode(frames_dir: str, prefix: str, fps: float, fmt: str, incl
     return args
 
 
-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):
+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
@@ -887,32 +960,68 @@ def step3_encode(frames_dir_state: str | None, prefix_state: str | None, orig_vi
         if detected:
             prefix = detected
     if not frames_dir or not prefix:
-        return None, None, "No frames available. Upload frames (ZIP/images) or run Step 1.", prog_html
+        yield None, "No frames available. Upload frames (ZIP/images) or run Step 1.", prog_html
+        return
 
     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)
 
+    # Build ffmpeg command
     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)
+
+    # Inject progress reporting
+    cmd.insert(1, "-progress")
+    cmd.insert(2, "pipe:2")
+
+    # Try to estimate total frames for progress %
+    total_frames = len(list(Path(frames_dir).glob(f"{prefix}_*.jpg"))) \
+                 + len(list(Path(frames_dir).glob(f"{prefix}_*.png")))
+
+    proc = subprocess.Popen(
+        cmd,
+        stderr=subprocess.PIPE,
+        stdout=subprocess.DEVNULL,
+        text=True,
+        bufsize=1,
+        cwd=frames_dir
+    )
 
     last_html = prog_html
+    current_frame = 0
+
     while True:
         line = proc.stderr.readline()
         if not line and proc.poll() is not None:
             break
-        if int(time.time()*10) % 5 == 0:
-            last_html = render_progress(50.0, "Encoding…")
-    ret = proc.wait()
 
+        if "frame=" in line:
+            try:
+                # parse `frame=123`
+                current_frame = int(line.strip().split("=")[-1])
+            except Exception:
+                pass
+
+            if total_frames > 0:
+                pct = min(100.0, (current_frame / total_frames) * 100.0)
+                last_html = render_progress(pct, f"Encoding… {current_frame}/{total_frames} frames")
+                yield None, f"Encoding in progress… {current_frame}/{total_frames}", last_html
+            else:
+                last_html = render_progress(50.0, "Encoding…")
+                yield None, "Encoding in progress…", last_html
+
+    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")
+        yield None, f"Encoding failed.\n\n{err}", last_html
+        return
+
+    yield str(out_file), f"Video created: {out_file.name}", render_progress(100.0, "Encoding complete")
+
 
 # ───────────────── Quick Mode — one click: All frames → Upscale ×4 → MP4 (audio)
 
@@ -1061,7 +1170,29 @@ def build_ui():
                     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>")
-
+                # Wire behavior: enable/disable param groups depending on mode / format
+                def _toggle_params(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_params,
+                    inputs=[mode, out_format],
+                    outputs=[every_seconds, nth_frame, exact_fps, jpg_quality, png_level],
+                )
+                out_format.change(
+                    _toggle_params,
+                    inputs=[mode, out_format],
+                    outputs=[every_seconds, nth_frame, exact_fps, jpg_quality, png_level],
+                )
+                # Initialize visibility
+                demo.load(_toggle_params, inputs=[mode, out_format], outputs=[every_seconds, nth_frame, exact_fps, jpg_quality, png_level])
+        
                 def update_estimate(vfile, mode_val, evs, nth, exfps, st, et):
                     if not vfile or not getattr(vfile, 'name', None):
                         return "Estimated output: —"