Update app.py

app.py

@@ -1,127 +1,157 @@
-import os
+import gradio as gr
 import cv2
 import numpy as np
+import os
 import tempfile
-from tqdm import tqdm
-import gradio as gr
-
+import subprocess
 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
+from basicsr.utils.download_util import load_file_from_url
+
+# --- Model Loading ---
+# We create a dictionary to cache models so they are only loaded once.
+model_cache = {}
 
+def get_upsampler(model_name='realesr-general-x4v3', denoise_strength=1):
+    """Loads and returns the specified RealESRGAN model."""
+    if model_name in model_cache:
+        return model_cache[model_name]
 
-# 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)
+        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')
+        file_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth'
+    else: # 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]
+        file_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
 
+    model_path = load_file_from_url(url=file_url, model_dir='weights', progress=True)
+    
     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
+        scale=netscale, model_path=model_path, model=model,
+        tile=64, tile_pad=10, pre_pad=10, half=False, gpu_id=None
     )
-
+    model_cache[model_name] = upsampler
     return upsampler
 
+def get_face_enhancer(upsampler, outscale):
+    """Loads and returns the GFPGAN face enhancer."""
+    key = 'face_enhancer'
+    if key in model_cache:
+        return model_cache[key]
+        
+    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
+    )
+    model_cache[key] = face_enhancer
+    return face_enhancer
+
+# --- Core Video Processing Function ---
+def enhance_video(video_path, model_name, denoise_strength, outscale, face_enhance, progress=gr.Progress(track_tqdm=True)):
+    """Enhances a video frame by frame and provides progress updates."""
+    if not video_path:
+        raise gr.Error("Please upload a video to enhance.")
+
+    try:
+        upsampler = get_upsampler(model_name, denoise_strength)
+        
+        face_enhancer = None
+        if face_enhance:
+            face_enhancer = get_face_enhancer(upsampler, outscale)
+
+        cap = cv2.VideoCapture(video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        
+        # Prepare output video writer
+        temp_dir = tempfile.mkdtemp()
+        enhanced_video_path = os.path.join(temp_dir, "enhanced_video.mp4")
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        writer = cv2.VideoWriter(enhanced_video_path, fourcc, fps, (width * outscale, height * outscale))
+
+        # Process each frame
+        for i in progress.tqdm(range(total_frames), desc="Enhancing Frames..."):
+            ret, frame = cap.read()
+            if not ret:
+                break
+            
+            if face_enhancer:
+                _, _, enhanced_frame = face_enhancer.enhance(frame, has_aligned=False, only_center_face=False, paste_back=True)
+            else:
+                enhanced_frame, _ = upsampler.enhance(frame, outscale=outscale)
+            
+            writer.write(enhanced_frame)
+        
+        cap.release()
+        writer.release()
+
+        # Merge audio back into the enhanced video
+        final_output_path = os.path.join(temp_dir, "final_output_with_audio.mp4")
+        audio_merge_cmd = f"ffmpeg -y -i {enhanced_video_path} -i {video_path} -c:v libx264 -crf 23 -preset fast -c:a aac -b:a 128k -map 0:v:0 -map 1:a:0 -shortest {final_output_path}"
+        subprocess.call(audio_merge_cmd, shell=True, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL)
+        
+        return final_output_path
+
+    except Exception as e:
+        print(traceback.format_exc())
+        raise gr.Error(f"An error occurred: {e}")
+
+# --- Gradio UI with Modern Design ---
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="violet"), title="🎥 AI Video Enhancer") as demo:
+    gr.Markdown(
+        """
+        # 🎥 AI Video Enhancer & Upscaler
+        Improve video quality, upscale resolution, and restore faces with cutting-edge AI. 
+        **Note:** Processing can be slow, especially for longer videos.
+        """
+    )
 
-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
+    with gr.Row(variant="panel"):
+        # --- Input Column ---
+        with gr.Column(scale=1):
+            video_input = gr.Video(label="🎬 Upload Your Video")
+            gr.Examples(
+                examples=["sample_video.mp4"], # Add path to your example video
+                inputs=[video_input],
+                label="Click an example to start"
+            )
+            
+            # Settings in a clean Accordion
+            with gr.Accordion("⚙️ Enhancement Options", open=True):
+                model_name = gr.Dropdown(
+                    choices=["realesr-general-x4v3", "RealESRGAN_x4plus_anime_6B"], 
+                    value="realesr-general-x4v3", 
+                    label="Model Type"
+                )
+                outscale = gr.Slider(1, 4, value=2, step=1, label="Upscale Factor")
+                face_enhance = gr.Checkbox(label="Restore Faces (GFPGAN)")
+                denoise_strength = gr.Slider(0, 1, value=0.5, step=0.1, label="Denoise Strength (for general model only)")
+            
+            enhance_btn = gr.Button("✨ Enhance Video", variant="primary")
+
+        # --- Output Column ---
+        with gr.Column(scale=1):
+            video_output = gr.Video(label="🌟 Enhanced Result")
+            download_file = gr.File(label="⬇️ Download Enhanced Video", visible=False)
+
+    # --- Event Logic ---
+    def on_submit(video, model, denoise, scale, face):
+        # When the button is clicked, start the enhancement and return the path to the output video.
+        # Also, make the download button visible.
+        output_path = enhance_video(video, model, denoise, scale, face)
+        return output_path, gr.update(value=output_path, visible=True)
+
+    enhance_btn.click(
+        fn=on_submit, 
+        inputs=[video_input, model_name, denoise_strength, outscale, face_enhance], 
+        outputs=[video_output, download_file]
+    )
 
-        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()
+if __name__ == '__main__':
+    demo.launch()