Upload 12 files

DEPLOY_HF.md

@@ -44,14 +44,15 @@ Do **not** upload `venv/`, `__pycache__/`, or model weights (downloaded at runti
 
 ## Troubleshooting
 - **Runtime `gradio.exceptions.Error: 'GPU task aborted'` / "ZeroGPU worker error"** —
-  the GPU worker was killed, usually because a **video** run was too big for one
-  120s GPU turn (e.g. many frames × 8K × per-frame **face-enhance**). The app now:
-  (a) **prefetches** imports+weights at startup; (b) caps **video** to ~4K output
-  with a **frame budget** (≈120 plain, 40 with face-enhance); and (c) has a
-  **time safety net** that stops before the limit and assembles the frames done
-  so far — so you get a (possibly shorter) video instead of an abort. For true 8K
-  video use a dedicated GPU; 8K is meant for the **Image** tab. Any remaining error
-  is printed to the logs via `traceback`.
+  the GPU worker was killed because a single call ran too long. **Video is now
+  processed in chunks**: the whole clip is split into batches
+  (≈{`chunk_frames_plain`} frames, {`chunk_frames_face`} with face-enhance), each
+  upscaled in its **own GPU call**, then joined with audio — so the full clip gets
+  done across multiple GPU turns instead of one short window. Output is capped to
+  ~4K; total length up to `max_total_video_seconds` (less with face-enhance). More
+  chunks use more of your ZeroGPU quota. For true 8K video use a dedicated GPU;
+  8K is meant for the **Image** tab. Startup **prefetches** imports+weights, and
+  any error is printed to the logs via `traceback`.
 - **Build fails building `basicsr` with `KeyError: '__version__'` / `get_version`** —
   the Space built on **Python 3.13**, whose PEP 667 broke BasicSR's `setup.py`
   (`exec()` + `locals()`). Fixed by pinning **`python_version: "3.10"`** in the

app.py

@@ -175,8 +175,21 @@ def enhance_image_fn(image, scale_mode, model_name, denoise, face_enhance,
 
 
 @spaces.GPU(duration=PROCESS_DURATION)
+def _gpu_upscale_batch(frame_paths, out_dir, out_start_index, outscale,
+                       model_name, denoise, face_enhance):
+    """One GPU call: upscale a batch of frames. Kept small enough to fit the
+    ZeroGPU time budget; the orchestrator calls this once per chunk."""
+    return pipeline.upscale_frames(
+        frame_paths, out_dir, int(out_start_index), float(outscale),
+        model_name, float(denoise), bool(face_enhance),
+    )
+
+
 def enhance_video_fn(video_file, scale_mode, model_name, denoise, face_enhance,
                      progress=gr.Progress()):
+    """Orchestrator (CPU): extract frames once, upscale them in batches — each
+    batch is a separate GPU call — then assemble. This processes the WHOLE clip
+    across multiple GPU turns instead of just one short window."""
     if video_file is None:
         gr.Warning("⚠️ Please upload a video first!")
         return None, "❌ No video uploaded"
@@ -184,20 +197,49 @@ def enhance_video_fn(video_file, scale_mode, model_name, denoise, face_enhance,
     def gp(v, t):
         progress(v, desc=t)
 
+    prep = None
     try:
         video_path = video_file if isinstance(video_file, str) else video_file.name
-        out_dir = tempfile.mkdtemp(prefix="enh_vid_")
-        out_path = os.path.join(out_dir, f"enhanced_{Path(video_path).stem}.mp4")
-        result, info = pipeline.enhance_video(
-            video_path, out_path, scale_mode=scale_mode, model_name=model_name,
-            denoise_strength=float(denoise), face_enhance=bool(face_enhance), progress=gp,
-        )
-        return result, info
+
+        # CPU: trim to overall cap + extract all frames.
+        prep = pipeline.prepare_video(video_path, scale_mode=scale_mode,
+                                      face_enhance=bool(face_enhance), progress=gp)
+        total = len(prep.files)
+        if total == 0:
+            return None, "❌ No frames found in the video"
+
+        batch = pipeline.batch_size(bool(face_enhance))
+        n_batches = (total + batch - 1) // batch
+
+        # GPU: process chunk by chunk (each a separate @spaces.GPU call).
+        out_idx = 0
+        for b, start in enumerate(range(0, total, batch)):
+            chunk = prep.files[start:start + batch]
+            end = min(start + batch, total)
+            gp(0.10 + 0.82 * (start / total),
+               f"🚀 Upscaling chunk {b + 1}/{n_batches} "
+               f"(frames {start + 1}–{end} of {total}) → {prep.ow}×{prep.oh}...")
+            out_idx = _gpu_upscale_batch(
+                chunk, prep.out_dir, out_idx, prep.outscale,
+                model_name, float(denoise), bool(face_enhance),
+            )
+
+        # CPU: join all upscaled frames + audio.
+        out_final = tempfile.mkdtemp(prefix="enh_vid_")
+        out_path = os.path.join(out_final, f"enhanced_{Path(video_path).stem}.mp4")
+        result, info = pipeline.finalize_video(prep, out_path, out_idx, progress=gp)
+        return result, info + f" · {n_batches} chunk(s)"
     except Exception as e:
         import traceback
         traceback.print_exc()   # surface the real error in the Space logs
         gr.Warning(f"Error: {e}")
         return None, f"❌ {e}"
+    finally:
+        if prep is not None:
+            try:
+                pipeline.cleanup(prep)
+            except Exception:
+                pass
 
 
 # ─── UI ────────────────────────────────────────────────────────
@@ -275,11 +317,13 @@ def create_app():
             ### Reaching 8K
             - **Images** upscale to 8K easily. Pick **Max → up to 8K** and the output is
               scaled to fill the 8K box (never larger). From ~1080p that's a 4× pass.
-            - **Video** is capped to **~4K output** and a **short frame budget** so a run
-              finishes inside one GPU turn (about {config.max_video_frames} frames plain,
-              **{config.max_video_frames_face} with face-enhance** since it's ~10× slower).
-              Longer input is auto-trimmed; if it still runs long it stops safely and
-              returns the frames done so far. **For true 8K video, use a dedicated GPU.**
+            - **Video** is processed **in chunks** — the whole clip is split into batches,
+              each upscaled in its own GPU turn, then joined back with audio. Output is
+              capped to **~4K**; total length up to **{config.max_total_video_seconds}s**
+              (**{config.max_total_video_seconds_face}s** with face-enhance, which is
+              ~10× slower). Longer input is trimmed. **For true 8K video, use a dedicated GPU.**
+            - More chunks = more GPU turns (and more of your ZeroGPU quota), but the whole
+              clip gets done instead of just the first second.
 
             ### Models
             - **General** — best for real-world photos/video; the **Denoise** dial trades grain vs. smoothness.

config.py

@@ -65,11 +65,14 @@ class AppConfig:
     max_video_duration: int = 8     # seconds (input; longer is auto-trimmed)
     max_output_long_edge: int = EIGHT_K_W   # cap output so it never exceeds 8K
 
-    # Video-specific caps so a run fits inside ONE ZeroGPU call:
+    # Video is processed in CHUNKS — each chunk is one @spaces.GPU call that
+    # fits the time budget; chunks are then joined. This lets the WHOLE clip be
+    # processed across multiple GPU calls instead of just one short window.
     max_video_long_edge: int = 3840          # video output capped to ~4K (8K stays for images)
-    max_video_frames: int = 120              # plain upscale frame budget
-    max_video_frames_face: int = 40          # face-enhance is ~10× slower per frame
-    safe_gpu_seconds: int = 90               # stop before ZeroGPU's ~120s limit & save partial
+    chunk_frames_plain: int = 120            # frames per GPU call (plain upscale)
+    chunk_frames_face: int = 30              # frames per GPU call (face-enhance is ~10× slower)
+    max_total_video_seconds: int = 20        # overall clip cap (plain); longer is trimmed
+    max_total_video_seconds_face: int = 10   # overall clip cap when face-enhance is on
 
     gpu_duration: int = 120         # seconds requested per @spaces.GPU call
     server_port: int = 7860

pipeline.py

@@ -2,20 +2,39 @@
 Enhancement pipeline.
 
 Image: decode -> Real-ESRGAN upscale -> encode.
-Video: (auto-trim if long) -> extract frames -> upscale each -> reassemble +audio.
+Video (CHUNKED so the whole clip is processed across multiple GPU calls):
+    prepare_video()  – CPU: (trim to total cap) + extract all frames
+    upscale_frames() – GPU: upscale ONE batch of frames  (called per chunk)
+    finalize_video() – CPU: assemble all upscaled frames + original audio
 """
 import os
 import cv2
 import time
 import numpy as np
 from pathlib import Path
-from typing import Optional, Callable, Tuple
+from dataclasses import dataclass
+from typing import Optional, Callable, Tuple, List
 
 from config import AppConfig, EIGHT_K_W, EIGHT_K_H
 from enhancer import RealESRGANEnhancer, compute_effective_scale
 from utils.video_processor import VideoProcessor
 
 
+@dataclass
+class VideoPrep:
+    temp_dir: str
+    frames_dir: str
+    out_dir: str
+    files: List[str]      # full paths to extracted source frames, in order
+    outscale: float
+    fps: float
+    width: int
+    height: int
+    ow: int
+    oh: int
+    trimmed_path: str     # (possibly trimmed) source, used for audio
+
+
 class EnhancePipeline:
     def __init__(self, config: AppConfig):
         self.config = config
@@ -43,30 +62,29 @@ class EnhancePipeline:
                f"{', face-enhanced' if face_enhance else ''})"
         return out, info
 
-    # ---------- video ----------
-    def enhance_video(
-        self, video_path: str, output_path: str, scale_mode: str = "4x",
-        model_name: str = "general", denoise_strength: float = 0.5,
-        face_enhance: bool = False, progress: Optional[Callable] = None,
-    ) -> Tuple[str, str]:
-        start_t = time.time()
+    # ---------- video: chunked ----------
+    def batch_size(self, face_enhance: bool) -> int:
+        return (getattr(self.config, "chunk_frames_face", 30) if face_enhance
+                else getattr(self.config, "chunk_frames_plain", 120))
+
+    def prepare_video(self, video_path: str, scale_mode: str = "4x",
+                      face_enhance: bool = False,
+                      progress: Optional[Callable] = None) -> VideoPrep:
+        """CPU: trim to the overall cap and extract all frames."""
         info = self.video_processor.get_video_info(video_path)
         fps = info.fps or 30.0
 
-        # Frame budget so ONE GPU call finishes in time. Face-enhance is far
-        # slower per frame, so it gets a much smaller budget.
-        budget = (getattr(self.config, "max_video_frames_face", 40) if face_enhance
-                  else getattr(self.config, "max_video_frames", 120))
-        dur_cap = getattr(self.config, "max_video_duration", 8)
-        target_seconds = min(dur_cap, budget / fps)
-        if info.duration and info.duration > target_seconds:
+        max_total = getattr(self.config, "max_total_video_seconds", 20)
+        if face_enhance:
+            max_total = min(max_total, getattr(self.config, "max_total_video_seconds_face", 10))
+        if info.duration and info.duration > max_total:
             self._p(progress, 0.03,
-                    f"✂️ Trimming to {target_seconds:.1f}s (~{int(target_seconds*fps)} frames) "
-                    f"to fit the GPU budget{' — face-enhance is slow' if face_enhance else ''}...")
-            video_path = self.video_processor.trim(video_path, 0, target_seconds)
+                    f"✂️ Trimming to {max_total}s (max total for "
+                    f"{'face-enhance' if face_enhance else 'this'} mode)...")
+            video_path = self.video_processor.trim(video_path, 0, max_total)
             info = self.video_processor.get_video_info(video_path)
 
-        # Effective scale, capped to the 8K box AND to the video resolution ceiling.
+        # Effective scale, capped to the 8K box AND the video resolution ceiling.
         outscale = compute_effective_scale(info.width, info.height, scale_mode,
                                            EIGHT_K_W, EIGHT_K_H)
         max_le = getattr(self.config, "max_video_long_edge", 3840)
@@ -77,56 +95,49 @@ class EnhancePipeline:
         temp_dir = os.path.join(self.config.video.temp_dir, "processing")
         frames_dir = os.path.join(temp_dir, "frames")
         out_dir = os.path.join(temp_dir, "out")
-        self.video_processor.cleanup_temp(temp_dir)   # clear any stale run
+        self.video_processor.cleanup_temp(temp_dir)
         os.makedirs(out_dir, exist_ok=True)
+
+        self._p(progress, 0.06, "🎞️ Extracting frames...")
+        self.video_processor.extract_frames_to_dir(video_path, frames_dir)
+        ext = self.config.video.frame_format
+        files = [os.path.join(frames_dir, f)
+                 for f in sorted(os.listdir(frames_dir)) if f.endswith(ext)]
+
+        return VideoPrep(temp_dir, frames_dir, out_dir, files, outscale, info.fps,
+                         info.width, info.height, ow, oh, video_path)
+
+    def upscale_frames(self, frame_paths: List[str], out_dir: str,
+                       out_start_index: int, outscale: float, model_name: str,
+                       denoise_strength: float, face_enhance: bool) -> int:
+        """GPU: upscale one batch of frames into out_dir, numbered contiguously
+        from out_start_index. Runs the model — call inside a GPU context.
+        Returns the next output index."""
+        self.enhancer.ensure_loaded(model_name, denoise_strength)
         ext = self.config.video.frame_format
-        safe_seconds = getattr(self.config, "safe_gpu_seconds", 90)
-
-        try:
-            self._p(progress, 0.08, "🎞️ Extracting frames...")
-            self.video_processor.extract_frames_to_dir(video_path, frames_dir)
-
-            self._p(progress, 0.12, f"Loading model ({model_name})...")
-            self.enhancer.ensure_loaded(model_name, denoise_strength, progress=progress)
-
-            files = sorted(f for f in os.listdir(frames_dir) if f.endswith(ext))
-            total = len(files)
-            self._p(progress, 0.15, f"Upscaling {total} frames ×{outscale:.2f} → {ow}×{oh}...")
-
-            processed = 0
-            stopped_early = False
-            for i, fname in enumerate(files):
-                # Safety net: stop before ZeroGPU's hard limit and keep what we have.
-                if time.time() - start_t > safe_seconds:
-                    stopped_early = True
-                    print(f"⏱️ GPU time budget ({safe_seconds}s) reached at frame {i}/{total}; "
-                          f"assembling the {processed} frames done so far.")
-                    break
-                frame = cv2.imread(os.path.join(frames_dir, fname), cv2.IMREAD_COLOR)
-                if frame is None:
-                    continue
-                up = self.enhancer.enhance(frame, outscale, face_enhance=face_enhance)
-                processed += 1
-                # Renumber sequentially so the output frames are always contiguous.
-                cv2.imwrite(os.path.join(out_dir, f"{processed:06d}.{ext}"), up)
-                if i % 2 == 0 or i == total - 1:
-                    p = 0.15 + (i / max(total, 1)) * 0.75
-                    self._p(progress, p, f"Upscaling frame {i+1}/{total}")
-
-            if processed == 0:
-                raise RuntimeError("No frames were processed.")
-
-            self._p(progress, 0.92, "🎬 Assembling video (+audio)...")
-            self.video_processor.assemble_video(out_dir, output_path, info.fps,
-                                                original_video=video_path)
-            elapsed = time.time() - start_t
-            note = "  ⚠️ partial (hit the GPU time limit)" if stopped_early else ""
-            msg = (f"✅ {info.width}×{info.height} → {ow}×{oh} (×{outscale:.2f}) · "
-                   f"{processed}/{total} frames · {elapsed:.0f}s{note}")
-            self._p(progress, 1.0, msg)
-            return output_path, msg
-        finally:
-            self.video_processor.cleanup_temp(temp_dir)
+        idx = out_start_index
+        for fp in frame_paths:
+            frame = cv2.imread(fp, cv2.IMREAD_COLOR)
+            if frame is None:
+                continue
+            up = self.enhancer.enhance(frame, outscale, face_enhance=face_enhance)
+            idx += 1
+            cv2.imwrite(os.path.join(out_dir, f"{idx:06d}.{ext}"), up)
+        return idx
+
+    def finalize_video(self, prep: VideoPrep, output_path: str, processed: int,
+                       progress: Optional[Callable] = None) -> Tuple[str, str]:
+        """CPU: assemble all upscaled frames + original audio."""
+        self._p(progress, 0.94, "🎬 Assembling video (+audio)...")
+        self.video_processor.assemble_video(prep.out_dir, output_path, prep.fps,
+                                            original_video=prep.trimmed_path)
+        msg = (f"✅ {prep.width}×{prep.height} → {prep.ow}×{prep.oh} "
+               f"(×{prep.outscale:.2f}) · {processed} frames")
+        self._p(progress, 1.0, msg)
+        return output_path, msg
+
+    def cleanup(self, prep: VideoPrep):
+        self.video_processor.cleanup_temp(prep.temp_dir)
 
     @staticmethod
     def _p(cb, v, t):