Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import os
 import tempfile
 import subprocess
+import traceback
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from realesrgan import RealESRGANer
 from realesrgan.archs.srvgg_arch import SRVGGNetCompact
@@ -14,7 +15,7 @@ from basicsr.utils.download_util import load_file_from_url
 # 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):
+def get_upsampler(model_name='realesr-general-x4v3'):
     """Loads and returns the specified RealESRGAN model."""
     if model_name in model_cache:
         return model_cache[model_name]
@@ -32,14 +33,14 @@ def get_upsampler(model_name='realesr-general-x4v3', denoise_strength=1):
     
     upsampler = RealESRGANer(
         scale=netscale, model_path=model_path, model=model,
-        tile=64, tile_pad=10, pre_pad=10, half=False, gpu_id=None
+        tile=64, tile_pad=10, pre_pad=10, half=True, gpu_id=None # Use half precision for speed
     )
     model_cache[model_name] = upsampler
     return upsampler
 
 def get_face_enhancer(upsampler, outscale):
     """Loads and returns the GFPGAN face enhancer."""
-    key = 'face_enhancer'
+    key = f'face_enhancer_{outscale}'
     if key in model_cache:
         return model_cache[key]
         
@@ -51,23 +52,23 @@ def get_face_enhancer(upsampler, outscale):
     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)):
+def enhance_video(video_path, model_name, 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)
+        upsampler = get_upsampler(model_name)
         
         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))
+        fps = cap.get(cv.CAP_PROP_FPS)
+        width = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
+        total_frames = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
         
         # Prepare output video writer
         temp_dir = tempfile.mkdtemp()
@@ -76,7 +77,7 @@ def enhance_video(video_path, model_name, denoise_strength, outscale, face_enhan
         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..."):
+        for _ in progress.tqdm(range(total_frames), desc="Enhancing Frames..."):
             ret, frame = cap.read()
             if not ret:
                 break
@@ -93,7 +94,7 @@ def enhance_video(video_path, model_name, denoise_strength, outscale, face_enhan
 
         # 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}"
+        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
@@ -106,46 +107,54 @@ def enhance_video(video_path, model_name, denoise_strength, outscale, face_enhan
 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.
         """
     )
 
-    with gr.Row(variant="panel"):
-        # --- Input Column ---
-        with gr.Column(scale=1):
-            video_input = gr.Video(label="🎬 Upload Your Video")
-          
-            # Settings in a clean Accordion
-            with gr.Accordion("⚙️ Enhancement Options", open=True):
+    # Top row for video previews
+    with gr.Row():
+        video_input = gr.Video(label="🎬 Original Video")
+        video_output = gr.Video(label="🌟 Enhanced Result")
+
+    # Panel for all settings and actions
+    with gr.Box():
+        with gr.Row():
+            # Left side for main settings
+            with gr.Column(scale=3):
                 model_name = gr.Dropdown(
                     choices=["realesr-general-x4v3", "RealESRGAN_x4plus_anime_6B"], 
                     value="realesr-general-x4v3", 
-                    label="Model Type"
+                    label="Model Type (General or Anime)"
                 )
                 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)
+            # Right side for the most important actions
+            with gr.Column(scale=1, min_width=200):
+                face_enhance = gr.Checkbox(label="✨ Restore Faces (GFPGAN)", value=False)
+                enhance_btn = gr.Button("🚀 Enhance Video", variant="primary")
+
+    # Examples and Download components
+    gr.Examples(
+        examples=["sample_video.mp4"], # Add path to your example video
+        inputs=[video_input],
+        label="Click an example to start"
+    )
+    download_file = gr.File(label="⬇️ Download Enhanced Video", visible=False)
 
     # --- Event Logic ---
-    def on_submit(video, model, denoise, scale, face):
+    def on_submit(video, model, 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)
+        output_path = enhance_video(video, model, 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], 
+        inputs=[video_input, model_name, outscale, face_enhance], 
         outputs=[video_output, download_file]
     )
 
 if __name__ == '__main__':
-    demo.launch()
+    demo.launch()