Upload app.py

app.py

@@ -0,0 +1,301 @@
+# app.py
+# AI Video Enhancer 4K - Optimized single GPU call with in-memory processing
+
+import os
+import shutil
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from typing import Tuple
+
+import gradio as gr
+import spaces
+import torch
+import numpy as np
+from PIL import Image
+import cv2
+from huggingface_hub import hf_hub_download
+from spandrel import ImageModelDescriptor, ModelLoader
+
+# Config
+TEMP_DIR = Path(tempfile.gettempdir()) / "hf_video_enhancer"
+TEMP_DIR.mkdir(parents=True, exist_ok=True)
+
+# Pre-download models at startup
+MODEL_PATHS = {}
+
+def ensure_model(scale: int) -> str:
+    """Download model weights (CPU/network only)."""
+    if scale not in MODEL_PATHS:
+        if scale == 2:
+            MODEL_PATHS[scale] = hf_hub_download(
+                repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x2.pth"
+            )
+        else:
+            MODEL_PATHS[scale] = hf_hub_download(
+                repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth"
+            )
+    return MODEL_PATHS[scale]
+
+try:
+    ensure_model(2)
+    ensure_model(4)
+    print("Models pre-downloaded successfully.")
+except Exception as e:
+    print(f"Model pre-download skipped: {e}")
+
+
+def run_cmd(cmd):
+    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    if p.returncode != 0:
+        raise RuntimeError(f"Command failed: {p.stderr.decode()}")
+    return p.stdout.decode()
+
+
+def probe_video(video_path: str) -> Tuple[float, int, int, float]:
+    cmd = [
+        "ffprobe", "-v", "error",
+        "-select_streams", "v:0",
+        "-show_entries", "stream=width,height,duration,r_frame_rate",
+        "-of", "default=noprint_wrappers=1:nokey=0",
+        video_path
+    ]
+    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    out = p.stdout.decode()
+    width = height = 0
+    duration = 0.0
+    fps = 30.0
+
+    for line in out.splitlines():
+        if line.startswith("width="):
+            width = int(line.split("=")[1])
+        elif line.startswith("height="):
+            height = int(line.split("=")[1])
+        elif line.startswith("duration="):
+            try:
+                duration = float(line.split("=")[1])
+            except:
+                pass
+        elif line.startswith("r_frame_rate="):
+            try:
+                fps_str = line.split("=")[1]
+                if "/" in fps_str:
+                    num, den = fps_str.split("/")
+                    fps = float(num) / float(den)
+                else:
+                    fps = float(fps_str)
+            except:
+                pass
+
+    return duration, width, height, fps
+
+
+def extract_frames(video_path: str, frames_dir: Path, max_frames: int = None):
+    frames_dir.mkdir(parents=True, exist_ok=True)
+    cmd = ["ffmpeg", "-y", "-i", video_path, "-vsync", "0"]
+    if max_frames:
+        cmd.extend(["-vframes", str(max_frames)])
+    cmd.append(str(frames_dir / "%06d.png"))
+    run_cmd(cmd)
+
+
+def reassemble_video(frames_dir: Path, audio_src: str, out_path: str, fps: float = 30.0):
+    tmp_video = str(frames_dir.parent / "tmp_video.mp4")
+    run_cmd([
+        "ffmpeg", "-y", "-framerate", str(fps),
+        "-i", str(frames_dir / "%06d.png"),
+        "-c:v", "libx264", "-preset", "veryfast", "-pix_fmt", "yuv420p",
+        "-crf", "18", tmp_video
+    ])
+
+    p = subprocess.run(
+        ["ffprobe", "-v", "error", "-select_streams", "a", "-show_entries",
+         "stream=codec_type", "-of", "default=noprint_wrappers=1", audio_src],
+        stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    )
+
+    if p.stdout.decode().strip():
+        run_cmd([
+            "ffmpeg", "-y", "-i", tmp_video, "-i", audio_src,
+            "-c:v", "copy", "-c:a", "aac",
+            "-map", "0:v:0", "-map", "1:a:0", out_path
+        ])
+        os.remove(tmp_video)
+    else:
+        shutil.move(tmp_video, out_path)
+
+
+def simple_upscale(img: np.ndarray, scale: int) -> np.ndarray:
+    h, w = img.shape[:2]
+    return cv2.resize(img, (w * scale, h * scale), interpolation=cv2.INTER_CUBIC)
+
+
+def load_frames_to_memory(frames_dir: Path) -> list:
+    """Load all frames into RAM (CPU work, not billed)."""
+    frame_files = sorted(frames_dir.glob("*.png"))
+    frames = []
+    for fp in frame_files:
+        img = cv2.imread(str(fp))
+        if img is not None:
+            frames.append(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+    return frames
+
+
+def save_frames_from_memory(frames: list, frames_dir: Path):
+    """Write enhanced frames back to disk (CPU work, not billed)."""
+    for idx, img_rgb in enumerate(frames):
+        out_path = frames_dir / f"{idx + 1:06d}.png"
+        cv2.imwrite(str(out_path), cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR))
+
+
+@spaces.GPU(duration=180)
+def enhance_all_frames_gpu(frames_rgb: list, model_path: str, scale: int = 4) -> list:
+    """
+    Enhance ALL frames in a SINGLE GPU call.
+    Frames are already in memory (no disk I/O here).
+    Uses torch.inference_mode() for faster inference with identical quality.
+    """
+    model = ModelLoader().load_from_file(model_path)
+    assert isinstance(model, ImageModelDescriptor)
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device).eval()
+
+    total = len(frames_rgb)
+    print(f"Model loaded on {device}, processing {total} frames in memory...")
+
+    enhanced = []
+    with torch.inference_mode():
+        for idx, img_rgb in enumerate(frames_rgb):
+            tensor = torch.from_numpy(img_rgb).permute(2, 0, 1).float().div(255.0)
+            tensor = tensor.unsqueeze(0).to(device)
+
+            output = model(tensor)
+            output = output.squeeze(0).cpu().clamp(0, 1).mul(255).byte()
+            output = output.permute(1, 2, 0).numpy()
+            enhanced.append(output)
+
+            if (idx + 1) % 10 == 0:
+                print(f"Processed {idx + 1}/{total}")
+
+    return enhanced
+
+
+def process_video(video_file, scale: int = 4, progress=gr.Progress()) -> Tuple[str, str]:
+    """Main video processing - single GPU call with in-memory frame processing."""
+    if video_file is None:
+        return "Please upload a video file.", None
+
+    ts = int(time.time() * 1000)
+    base_dir = TEMP_DIR / f"job_{ts}"
+    base_dir.mkdir(parents=True, exist_ok=True)
+    in_path = base_dir / "input_video"
+
+    try:
+        shutil.copy(video_file, in_path)
+    except Exception as e:
+        return f"Error: {e}", None
+
+    try:
+        duration, w, h, fps = probe_video(str(in_path))
+    except Exception as e:
+        shutil.rmtree(base_dir, ignore_errors=True)
+        return f"Error probing video: {e}", None
+
+    if duration <= 0:
+        shutil.rmtree(base_dir, ignore_errors=True)
+        return "Could not determine video duration.", None
+
+    max_seconds = 10
+    max_frames = int(fps * max_seconds)
+
+    progress(0.05, f"Video: {w}x{h} @ {fps:.1f}fps, extracting up to {max_seconds}s...")
+
+    frames_dir = base_dir / "frames"
+    try:
+        extract_frames(str(in_path), frames_dir, max_frames)
+    except Exception as e:
+        shutil.rmtree(base_dir, ignore_errors=True)
+        return f"Failed extracting frames: {e}", None
+
+    num_frames = len(list(frames_dir.glob("*.png")))
+    progress(0.15, f"Loading {num_frames} frames into memory...")
+
+    # Load all frames into RAM (CPU work, no GPU needed)
+    frames_rgb = load_frames_to_memory(frames_dir)
+    if not frames_rgb:
+        shutil.rmtree(base_dir, ignore_errors=True)
+        return "No frames extracted.", None
+
+    # Ensure model weights are cached (CPU/network only)
+    progress(0.20, "Preparing model...")
+    model_path = ensure_model(scale)
+
+    progress(0.25, f"Enhancing {len(frames_rgb)} frames on GPU (single call)...")
+
+    use_fallback = False
+    try:
+        # === THE SINGLE GPU CALL ===
+        enhanced_frames = enhance_all_frames_gpu(frames_rgb, model_path, scale)
+        print(f"Enhanced {len(enhanced_frames)} frames with Real-ESRGAN")
+    except Exception as e:
+        print(f"GPU enhancement failed: {e}")
+        print("Using fallback bicubic upscaling...")
+        use_fallback = True
+        enhanced_frames = [simple_upscale(f, scale) for f in frames_rgb]
+
+    progress(0.80, "Writing enhanced frames...")
+    save_frames_from_memory(enhanced_frames, frames_dir)
+
+    # Free memory
+    del frames_rgb, enhanced_frames
+
+    progress(0.85, "Reassembling video...")
+    out_video = base_dir / "enhanced_output.mp4"
+    try:
+        reassemble_video(frames_dir, str(in_path), str(out_video), fps)
+    except Exception as e:
+        shutil.rmtree(base_dir, ignore_errors=True)
+        return f"Failed reassembling: {e}", None
+
+    shutil.rmtree(frames_dir, ignore_errors=True)
+
+    try:
+        _, out_w, out_h, _ = probe_video(str(out_video))
+        method = "bicubic" if use_fallback else "Real-ESRGAN"
+        progress(1.0, "Done!")
+        return f"Done: {w}x{h} -> {out_w}x{out_h} ({method}, {num_frames} frames)", str(out_video)
+    except:
+        return "Done!", str(out_video)
+
+
+# Gradio UI
+with gr.Blocks(title="AI Video Enhancer", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# AI Video Enhancer")
+    gr.Markdown("Upscale videos using Real-ESRGAN AI. **Log in above for more GPU quota!**")
+
+    # LOGIN BUTTON - This allows ZeroGPU to recognize your Pro account
+    gr.LoginButton()
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            video_in = gr.File(
+                label="Upload video (max 10 sec processed)",
+                file_types=[".mp4", ".avi", ".mov", ".mkv", ".webm"]
+            )
+            scale_choice = gr.Radio(choices=[2, 4], value=4, label="Upscale Factor")
+            btn = gr.Button("Enhance", variant="primary")
+            status = gr.Textbox(label="Status", interactive=False)
+        with gr.Column(scale=1):
+            out_video = gr.Video(label="Result")
+
+    gr.Markdown(
+        "**Limit: 10 seconds of video** (ZeroGPU quota). "
+        "Log in to HuggingFace for more!"
+    )
+
+    btn.click(fn=process_video, inputs=[video_in, scale_choice], outputs=[status, out_video])
+
+if __name__ == "__main__":
+    demo.launch()