# app.py # AI Video Enhancer with OAuth login for ZeroGPU quota import os import shutil import subprocess import tempfile import time from pathlib import Path import gradio as gr import spaces import torch import numpy as np from PIL import Image # ============== Configuration ============== TEMP_DIR = Path(tempfile.gettempdir()) / "video_enhancer" TEMP_DIR.mkdir(parents=True, exist_ok=True) MODELS = { 2: "caidas/swin2SR-classical-sr-x2-64", 4: "caidas/swin2SR-classical-sr-x4-48", } # ============== Video Utilities ============== def run_ffmpeg(cmd: list) -> str: result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: raise RuntimeError(f"FFmpeg error: {result.stderr}") return result.stdout def get_video_info(path: str) -> dict: import json cmd = [ "ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", "-show_format", path ] result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: raise RuntimeError("Could not read video info") data = json.loads(result.stdout) video_stream = next((s for s in data.get("streams", []) if s["codec_type"] == "video"), None) if not video_stream: raise RuntimeError("No video stream found") fps_str = video_stream.get("r_frame_rate", "30/1") if "/" in fps_str: num, den = map(float, fps_str.split("/")) fps = num / den if den != 0 else 30.0 else: fps = float(fps_str) return { "width": int(video_stream.get("width", 0)), "height": int(video_stream.get("height", 0)), "fps": fps, "duration": float(data.get("format", {}).get("duration", 0)), "has_audio": any(s["codec_type"] == "audio" for s in data.get("streams", [])) } def extract_frames(video_path: str, output_dir: Path, max_frames: int = None) -> int: output_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(output_dir / "frame_%06d.png")) run_ffmpeg(cmd) return len(list(output_dir.glob("*.png"))) def create_video(frames_dir: Path, output_path: str, fps: float, audio_source: str = None): temp_video = str(Path(output_path).parent / "temp_no_audio.mp4") run_ffmpeg([ "ffmpeg", "-y", "-framerate", str(fps), "-i", str(frames_dir / "frame_%06d.png"), "-c:v", "libx264", "-preset", "medium", "-crf", "18", "-pix_fmt", "yuv420p", temp_video ]) if audio_source: try: info = get_video_info(audio_source) if info["has_audio"]: run_ffmpeg([ "ffmpeg", "-y", "-i", temp_video, "-i", audio_source, "-c:v", "copy", "-c:a", "aac", "-map", "0:v:0", "-map", "1:a:0", output_path ]) os.remove(temp_video) return except: pass shutil.move(temp_video, output_path) # ============== Image Enhancement ============== @spaces.GPU(duration=60) def enhance_batch(images: list, scale: int = 2) -> list: """Enhance a batch of images using Swin2SR.""" from transformers import Swin2SRForImageSuperResolution, Swin2SRImageProcessor model_id = MODELS[scale] processor = Swin2SRImageProcessor.from_pretrained(model_id) model = Swin2SRForImageSuperResolution.from_pretrained(model_id) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device).eval() enhanced = [] for img in images: inputs = processor(img, return_tensors="pt").to(device) with torch.no_grad(): outputs = model(**inputs) output = outputs.reconstruction.squeeze().cpu().clamp(0, 1) output = output.permute(1, 2, 0).numpy() output = (output * 255).astype(np.uint8) enhanced.append(Image.fromarray(output)) return enhanced # ============== Main Processing ============== def process_video(video_file, scale: int, profile: gr.OAuthProfile | None, progress=gr.Progress()): """Main video processing function""" # Show login status user_msg = f"Processing as: {profile.name}" if profile else "Not logged in (limited quota)" print(user_msg) if video_file is None: return "⚠️ Please upload a video file.", None job_id = f"job_{int(time.time() * 1000)}" work_dir = TEMP_DIR / job_id work_dir.mkdir(parents=True, exist_ok=True) input_path = work_dir / "input_video" frames_dir = work_dir / "frames" output_path = work_dir / "output.mp4" try: shutil.copy(video_file, input_path) progress(0.05, "Analyzing video...") info = get_video_info(str(input_path)) max_frames = min(int(info["fps"] * 10), 300) progress(0.1, "Extracting frames...") num_frames = extract_frames(str(input_path), frames_dir, max_frames) if num_frames == 0: return "❌ Could not extract frames from video.", None progress(0.15, f"Processing {num_frames} frames...") frame_files = sorted(frames_dir.glob("*.png")) batch_size = 4 for batch_start in range(0, len(frame_files), batch_size): batch_end = min(batch_start + batch_size, len(frame_files)) batch_files = frame_files[batch_start:batch_end] batch_images = [Image.open(f).convert("RGB") for f in batch_files] try: enhanced_images = enhance_batch(batch_images, scale) for img, path in zip(enhanced_images, batch_files): img.save(path) except Exception as e: print(f"GPU enhancement failed: {e}, using fallback...") for img, path in zip(batch_images, batch_files): new_size = (img.width * scale, img.height * scale) img.resize(new_size, Image.LANCZOS).save(path) pct = 0.15 + 0.75 * (batch_end / len(frame_files)) progress(pct, f"Enhanced {batch_end}/{len(frame_files)} frames") progress(0.92, "Creating video...") create_video(frames_dir, str(output_path), info["fps"], str(input_path)) shutil.rmtree(frames_dir, ignore_errors=True) out_info = get_video_info(str(output_path)) progress(1.0, "Done!") return ( f"✅ Enhanced: {info['width']}x{info['height']} → " f"{out_info['width']}x{out_info['height']} ({scale}x)", str(output_path) ) except Exception as e: shutil.rmtree(work_dir, ignore_errors=True) return f"❌ Error: {str(e)}", None def greet(profile: gr.OAuthProfile | None) -> str: if profile is None: return "👋 **Not logged in** - Click 'Sign in with Hugging Face' above for full GPU quota!" return f"👋 Welcome **{profile.name}**! You have Pro GPU quota." # ============== Gradio Interface ============== with gr.Blocks(title="AI Video Enhancer", theme=gr.themes.Soft()) as demo: gr.Markdown("# 🎬 AI Video Enhancer") # Login button and status gr.LoginButton() login_status = gr.Markdown() demo.load(greet, inputs=None, outputs=login_status) gr.Markdown("---") with gr.Row(): with gr.Column(scale=2): video_input = gr.File( label="📁 Upload Video", file_types=[".mp4", ".avi", ".mov", ".mkv", ".webm"] ) scale_selector = gr.Radio( choices=[2, 4], value=2, label="🔍 Upscale Factor" ) enhance_btn = gr.Button("🚀 Enhance Video", variant="primary", size="lg") status_text = gr.Textbox(label="Status", interactive=False) with gr.Column(scale=1): video_output = gr.Video(label="🎥 Enhanced Video") enhance_btn.click( fn=process_video, inputs=[video_input, scale_selector], outputs=[status_text, video_output] ) if __name__ == "__main__": demo.launch()