Upload 2 files

app.py

@@ -0,0 +1,280 @@
+# app.py
+# AI Video Enhancer 4K - Gradio app for Hugging Face Spaces
+# MCP-enabled version
+
+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
+
+# Config
+TEMP_DIR = Path(tempfile.gettempdir()) / "hf_video_enhancer"
+TEMP_DIR.mkdir(parents=True, exist_ok=True)
+
+
+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):
+    frames_dir.mkdir(parents=True, exist_ok=True)
+    run_cmd([
+        "ffmpeg", "-y", "-i", video_path,
+        "-vsync", "0",
+        str(frames_dir / "%06d.png")
+    ])
+
+
+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:
+    """Simple bicubic upscaling using OpenCV"""
+    h, w = img.shape[:2]
+    return cv2.resize(img, (w * scale, h * scale), interpolation=cv2.INTER_CUBIC)
+
+
+@spaces.GPU(duration=120)
+def enhance_with_realesrgan(frames_dir: str, scale: int = 4) -> int:
+    """
+    Enhance frames using Real-ESRGAN via Spandrel.
+    Separated function with GPU decorator for cleaner ZeroGPU handling.
+    """
+    from spandrel import ImageModelDescriptor, ModelLoader
+
+    frames_path = Path(frames_dir)
+    frame_files = sorted(frames_path.glob("*.png"))
+    total = len(frame_files)
+
+    if total == 0:
+        return 0
+
+    if scale == 2:
+        model_path = hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x2.pth")
+    else:
+        model_path = hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth")
+
+    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()
+
+    print(f"Model loaded on {device}, processing {total} frames...")
+
+    for idx, frame_path in enumerate(frame_files):
+        img = cv2.imread(str(frame_path))
+        if img is None:
+            continue
+
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        tensor = torch.from_numpy(img_rgb).permute(2, 0, 1).float().div(255.0)
+        tensor = tensor.unsqueeze(0).to(device)
+
+        with torch.no_grad():
+            output = model(tensor)
+
+        output = output.squeeze(0).cpu().clamp(0, 1).mul(255).byte()
+        output = output.permute(1, 2, 0).numpy()
+        output_bgr = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+
+        cv2.imwrite(str(frame_path), output_bgr)
+
+        if (idx + 1) % 5 == 0:
+            print(f"Processed {idx + 1}/{total}")
+
+    return total
+
+
+def process_video(video_file: str, scale: int = 4) -> Tuple[str, str]:
+    """
+    Upscale and enhance a video using Real-ESRGAN AI super-resolution.
+
+    Takes a low-resolution or standard video and outputs a sharper, higher-resolution
+    version using the Real-ESRGAN model. Audio is preserved. Processing is limited
+    to approximately 30 seconds of video due to GPU constraints.
+
+    Args:
+        video_file: Path to the input video file. Supported formats: mp4, avi, mov, mkv, webm.
+        scale: Upscaling factor. Use 2 for 2x resolution boost or 4 for 4x (default: 4).
+
+    Returns:
+        A tuple of (status_message, output_video_path). The status message describes
+        the resolution change (e.g. '480x270 → 1920x1080'). The output_video_path
+        is the local path to the enhanced MP4 file.
+    """
+    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_frames = int(fps * 30)
+
+    print(f"Video: {w}x{h}, {duration:.1f}s, {fps:.1f}fps")
+
+    frames_dir = base_dir / "frames"
+    try:
+        extract_frames(str(in_path), frames_dir)
+    except Exception as e:
+        shutil.rmtree(base_dir, ignore_errors=True)
+        return f"Failed extracting frames: {e}", None
+
+    frame_files = sorted(frames_dir.glob("*.png"))
+    num_frames = len(frame_files)
+
+    if num_frames > max_frames:
+        print(f"Limiting from {num_frames} to {max_frames} frames")
+        for f in frame_files[max_frames:]:
+            f.unlink()
+        num_frames = max_frames
+
+    print(f"Processing {num_frames} frames...")
+
+    try:
+        enhanced = enhance_with_realesrgan(str(frames_dir), scale)
+        print(f"Enhanced {enhanced} frames")
+    except Exception as e:
+        print(f"Enhancement failed: {e}")
+        print("Using fallback bicubic upscaling...")
+        try:
+            for fp in sorted(frames_dir.glob("*.png")):
+                img = cv2.imread(str(fp))
+                if img is not None:
+                    upscaled = simple_upscale(img, scale)
+                    cv2.imwrite(str(fp), upscaled)
+        except Exception as e2:
+            shutil.rmtree(base_dir, ignore_errors=True)
+            return f"Enhancement failed: {e}", None
+
+    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))
+        return f"✅ Done! {w}x{h} → {out_w}x{out_h}", 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 enhancement.")
+
+    gr.LoginButton()
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            video_in = gr.File(label="Upload video", 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("**Note:** Limited to ~30 seconds for ZeroGPU. Longer videos will be truncated.")
+
+    btn.click(fn=process_video, inputs=[video_in, scale_choice], outputs=[status, out_video])
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

requirements.txt

@@ -0,0 +1,8 @@
+gradio[mcp]
+spaces
+torch
+numpy
+Pillow
+opencv-python-headless
+huggingface_hub
+spandrel