Create app.py

app.py

@@ -0,0 +1,127 @@
+import os
+import cv2
+import numpy as np
+import tempfile
+from tqdm import tqdm
+import gradio as gr
+
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.download_util import load_file_from_url
+from realesrgan import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+from gfpgan import GFPGANer
+
+
+# Load models
+def load_model(model_name, denoise_strength=1.0):
+    if model_name == 'RealESRGAN_x4plus_anime_6B':
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
+                        num_block=6, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
+    elif model_name == 'realesr-general-x4v3':
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64,
+                                num_conv=32, upscale=4, act_type='prelu')
+        netscale = 4
+        file_url = [
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
+        ]
+
+    model_path = os.path.join('weights', model_name + '.pth')
+    os.makedirs('weights', exist_ok=True)
+
+    if not os.path.isfile(model_path):
+        for url in file_url:
+            model_path = load_file_from_url(url=url, model_dir='weights', progress=True)
+
+    dni_weight = None
+    if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
+        model_path = [
+            os.path.join('weights', 'realesr-general-x4v3.pth'),
+            os.path.join('weights', 'realesr-general-wdn-x4v3.pth')
+        ]
+        dni_weight = [denoise_strength, 1 - denoise_strength]
+
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=128,
+        tile_pad=10,
+        pre_pad=10,
+        half=False,
+        gpu_id=None
+    )
+
+    return upsampler
+
+
+def enhance_video(video_path, model_name, denoise_strength, face_enhance, outscale):
+    upsampler = load_model(model_name, denoise_strength)
+
+    if face_enhance:
+        face_enhancer = GFPGANer(
+            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
+            upscale=outscale,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=upsampler
+        )
+
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    temp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
+    out_path = temp_out.name
+
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    writer = cv2.VideoWriter(out_path, fourcc, fps, (w * outscale, h * outscale))
+
+    for _ in tqdm(range(total_frames), desc="Enhancing video"):
+        success, frame = cap.read()
+        if not success:
+            break
+
+        try:
+            if face_enhance:
+                _, _, enhanced = face_enhancer.enhance(frame, has_aligned=False, only_center_face=False, paste_back=True)
+            else:
+                enhanced, _ = upsampler.enhance(frame, outscale=outscale)
+            writer.write(enhanced)
+        except RuntimeError as e:
+            print("Runtime error:", e)
+            continue
+
+    cap.release()
+    writer.release()
+
+    return out_path
+
+
+def gradio_interface(video, model_name, denoise_strength, face_enhance, outscale):
+    if video is None:
+        return None
+    return enhance_video(video, model_name, denoise_strength, face_enhance, outscale)
+
+
+demo = gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.Video(label="Upload a short video (<30s)"),
+        gr.Dropdown(["realesr-general-x4v3", "RealESRGAN_x4plus_anime_6B"], label="Model", value="realesr-general-x4v3"),
+        gr.Slider(0, 1, step=0.1, value=1.0, label="Denoise Strength"),
+        gr.Checkbox(label="Enable Face Enhancement (GFPGAN)", value=False),
+        gr.Slider(1, 4, step=1, value=2, label="Upscale Factor")
+    ],
+    outputs=gr.Video(label="Enhanced Video Output"),
+    title="🎬 AI Video Enhancer",
+    description="Upscale your videos with Real-ESRGAN and optional face enhancement using GFPGAN. Optimized for Hugging Face CPU Spaces."
+)
+
+if __name__ == "__main__":
+    demo.launch()