Create app.py

app.py

@@ -0,0 +1,293 @@
+# app.py
+# AI Video Enhancer 4K - Gradio app for Hugging Face Spaces
+# NOTE: This app attempts to use GFPGAN and Real-ESRGAN if installed.
+# If they're not available (common on CPU-only environments), it falls back to ffmpeg-based upscale.
+
+import os
+import shutil
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from typing import Tuple
+
+import gradio as gr
+from PIL import Image
+
+# Try to import optional enhancement libs
+try:
+    import torch
+    from gfpgan import GFPGANer  # type: ignore
+    from realesrgan import RealESRGAN  # type: ignore
+    HAVE_ENHANCERS = True
+except Exception:
+    HAVE_ENHANCERS = False
+
+# Config
+MAX_SECONDS = 30          # maximum video length
+TEMP_DIR = Path(tempfile.gettempdir()) / "hf_video_enhancer"
+TEMP_DIR.mkdir(parents=True, exist_ok=True)
+
+def run_cmd(cmd):
+    """Run shell command, raise if failed."""
+    p = subprocess.run(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    if p.returncode != 0:
+        raise RuntimeError(f"Command failed: {' '.join(cmd)}\nSTDOUT:{p.stdout.decode()}\nSTDERR:{p.stderr.decode()}")
+    return p.stdout.decode()
+
+def probe_video(video_path: str) -> Tuple[float, int, int]:
+    """Return (duration_seconds, width, height) using ffprobe."""
+    cmd = [
+        "ffprobe", "-v", "error",
+        "-select_streams", "v:0",
+        "-show_entries", "stream=width,height,duration",
+        "-of", "default=noprint_wrappers=1:nokey=0",
+        video_path
+    ]
+    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    out = p.stdout.decode()
+    # parse
+    width = height = 0
+    duration = 0.0
+    for line in out.splitlines():
+        if line.startswith("width="):
+            width = int(line.split("=", 1)[1])
+        if line.startswith("height="):
+            height = int(line.split("=", 1)[1])
+        if line.startswith("duration="):
+            try:
+                duration = float(line.split("=", 1)[1])
+            except:
+                duration = 0.0
+    return duration, width, height
+
+def extract_frames(video_path: str, frames_dir: Path):
+    """Extract frames as PNG to frames_dir."""
+    frames_dir.mkdir(parents=True, exist_ok=True)
+    # %06d.png
+    cmd = [
+        "ffmpeg", "-y", "-i", video_path,
+        "-vsync", "0",
+        str(frames_dir / "%06d.png")
+    ]
+    run_cmd(cmd)
+
+def reassemble_video(frames_dir: Path, audio_src: str, out_path: str, fps: float = 30.0):
+    """Reassemble frames to video and add original audio (if exists)."""
+    # Determine fps if possible by probing source
+    cmd_probe = [
+        "ffprobe", "-v", "error", "-select_streams", "v:0",
+        "-show_entries", "stream=r_frame_rate",
+        "-of", "default=noprint_wrappers=1:nokey=1", audio_src
+    ]
+    try:
+        p = subprocess.run(cmd_probe, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        r = p.stdout.decode().strip()
+        if "/" in r:
+            a, b = r.split("/")
+            fps = float(a) / float(b)
+    except Exception:
+        pass
+
+    # Create video from frames
+    tmp_video = str(frames_dir.parent / "tmp_video_no_audio.mp4")
+    cmd_encode = [
+        "ffmpeg", "-y",
+        "-framerate", str(fps),
+        "-i", str(frames_dir / "%06d.png"),
+        "-c:v", "libx264", "-preset", "veryfast", "-pix_fmt", "yuv420p",
+        tmp_video
+    ]
+    run_cmd(cmd_encode)
+
+    # Add audio if present
+    # First check if the source has an audio stream
+    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)
+    has_audio = bool(p.stdout.decode().strip())
+    if has_audio:
+        cmd_mux = [
+            "ffmpeg", "-y",
+            "-i", tmp_video,
+            "-i", audio_src,
+            "-c:v", "copy",
+            "-c:a", "aac",
+            "-map", "0:v:0",
+            "-map", "1:a:0",
+            out_path
+        ]
+        run_cmd(cmd_mux)
+        os.remove(tmp_video)
+    else:
+        # just rename
+        shutil.move(tmp_video, out_path)
+
+def simple_upscale_with_ffmpeg(frames_dir: Path, scale_factor: int = 2):
+    """Basic ffmpeg-based upscale (nearest/linear) as a fallback."""
+    # iterate frames and upscale in-place
+    for p in sorted(frames_dir.glob("*.png")):
+        tmp = str(p) + ".tmp.png"
+        cmd = [
+            "ffmpeg", "-y", "-i", str(p),
+            "-vf", f"scale=iw*{scale_factor}:ih*{scale_factor}:flags=lanczos",
+            tmp
+        ]
+        run_cmd(cmd)
+        os.replace(tmp, p)
+
+def load_enhancers(device="cuda"):
+    """Initialize GFPGAN and Real-ESRGAN if available. Returns (gfpganer, realesrgan) or (None, None)."""
+    if not HAVE_ENHANCERS:
+        return None, None
+    # GFPGAN
+    try:
+        # GFPGANer parameters depend on package version; this is best-effort
+        gfpganer = GFPGANer(
+            model_path=None,  # let package find default
+            upscale=1,
+            arch='clean', 
+            channel_multiplier=2,
+            bg_upsampler=None
+        )
+    except Exception:
+        gfpganer = None
+
+    try:
+        # RealESRGAN uses a model name typically; attempt to load default x2 model
+        realesrgan = RealESRGAN(device, scale=4)
+        realesrgan.load_weights('RealESRGAN_x4plus')  # may raise if not present
+    except Exception:
+        realesrgan = None
+
+    return gfpganer, realesrgan
+
+def enhance_frames(frames_dir: Path, progress= None):
+    """Apply GFPGAN (face restore) then Real-ESRGAN upscaling if available. Otherwise fallback to ffmpeg upscale."""
+    if HAVE_ENHANCERS:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        gfpganer, realesrgan = load_enhancers(device)
+        frame_paths = sorted(frames_dir.glob("*.png"))
+        n = len(frame_paths)
+        for i, fp in enumerate(frame_paths, start=1):
+            img = Image.open(fp).convert("RGB")
+            # face restore
+            try:
+                if gfpganer is not None:
+                    _, restored = gfpganer.enhance(np.array(img), has_aligned=False, only_center_face=False, paste_back=True)
+                    img = Image.fromarray(restored)
+            except Exception:
+                pass
+            # upscale
+            try:
+                if realesrgan is not None:
+                    with torch.no_grad():
+                        out = realesrgan.predict(img)
+                        out.save(fp)
+                else:
+                    # fallback to ffmpeg scaling for this frame
+                    tmp = str(fp) + ".tmp.png"
+                    cmd = [
+                        "ffmpeg", "-y", "-i", str(fp),
+                        "-vf", "scale=iw*2:ih*2:flags=lanczos",
+                        tmp
+                    ]
+                    run_cmd(cmd)
+                    os.replace(tmp, str(fp))
+            except Exception:
+                # final fallback: leave the frame as-is
+                pass
+            if progress:
+                progress(i / n)
+    else:
+        # simple ffmpeg upscale 2x
+        simple_upscale_with_ffmpeg(frames_dir, scale_factor=2)
+
+# The main processing pipeline
+def process_video(video_file) -> Tuple[str, str]:
+    """
+    Accepts an uploaded video file from Gradio,
+    processes it and returns (message, path_to_result_video)
+    """
+    # Save upload to temp
+    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"
+    with open(in_path, "wb") as f:
+        f.write(video_file.read())
+
+    # probe
+    try:
+        duration, w, h = probe_video(str(in_path))
+    except Exception as e:
+        shutil.rmtree(base_dir, ignore_errors=True)
+        return f"Error probing video: {e}", ""
+
+    if duration > MAX_SECONDS:
+        shutil.rmtree(base_dir, ignore_errors=True)
+        return f"Video too long: {duration:.1f}s (limit {MAX_SECONDS}s). Trim it and try again.", ""
+
+    # extract frames
+    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}", ""
+
+    # enhancement step
+    # Using a progress hook from Gradio is a bit awkward; we'll just run and hope for the best
+    try:
+        enhance_frames(frames_dir)
+    except Exception as e:
+        # continue, fallback allowed
+        print(f"Enhancement failed: {e}")
+
+    # reassemble and add audio
+    out_video = base_dir / "enhanced_output.mp4"
+    try:
+        reassemble_video(frames_dir, str(in_path), str(out_video))
+    except Exception as e:
+        shutil.rmtree(base_dir, ignore_errors=True)
+        return f"Failed to reassemble video: {e}", ""
+
+    # Optionally: cleanup frames to save space
+    try:
+        shutil.rmtree(frames_dir)
+    except Exception:
+        pass
+
+    # Serve out_video path as string (Gradio will handle file serving)
+    return "Processing complete. Download below.", str(out_video)
+
+# Gradio UI
+with gr.Blocks(title="AI Video Enhancer 4K") as demo:
+    gr.Markdown("# AI Video Enhancer 4K")
+    gr.Markdown("Upload a short video (<= 30s). The app will attempt to enhance faces and upscale frames. Heavy models may not run on CPU-only free environments.")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            video_in = gr.File(label="Upload video (mp4/avi/mov)", file_count="single")
+            btn = gr.Button("Enhance Video")
+            status = gr.Textbox(label="Status", interactive=False)
+        with gr.Column(scale=1):
+            out_video = gr.Video(label="Enhanced video")
+            download_btn = gr.Button("Download enhanced video")
+
+    def on_click_process(file_obj):
+        if not file_obj:
+            return "Please upload a video file.", None
+        try:
+            msg, path = process_video(file_obj)
+            if path:
+                return msg, path
+            else:
+                return msg, None
+        except Exception as e:
+            return f"Unexpected error: {e}", None
+
+    btn.click(fn=on_click_process, inputs=[video_in], outputs=[status, out_video])
+
+    gr.Markdown("**Notes:** If the Space is CPU-only or lacks the heavy model weights, the app will run a simpler ffmpeg upscale. For realistic Real-ESRGAN/GFPGAN results, enable GPU on the Space and ensure model weights are installed in `./weights` or accessible to the packages used.")
+
+if __name__ == "__main__":
+    demo.launch()