Add Gradio interface for Colab

google_colabs/gradio_app.py

@@ -0,0 +1,242 @@
+"""
+Gradio Interface for Video Subtitle Remover
+Optimized for Google Colab
+"""
+
+import gradio as gr
+import cv2
+import os
+import sys
+import tempfile
+from pathlib import Path
+
+# Detect if running in Colab
+IN_COLAB = 'google.colab' in sys.modules
+
+# Set correct paths for Colab
+if IN_COLAB:
+    PROJECT_ROOT = '/content/sub-remover'
+else:
+    PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+# Add project to path
+sys.path.insert(0, PROJECT_ROOT)
+sys.path.insert(0, os.path.join(PROJECT_ROOT, 'backend'))
+
+from backend import config
+from backend.config import InpaintMode
+from backend.main import SubtitleRemover
+
+
+def process_video(
+    video_file,
+    algorithm,
+    sttn_skip_detection,
+    sttn_max_load,
+    lama_super_fast,
+    sd_steps,
+    sd_guidance,
+    diffueraser_steps,
+    diffueraser_max_load,
+    progress=gr.Progress()
+):
+    """Process video to remove subtitles"""
+    if video_file is None:
+        return None, "❌ Please upload a video first"
+    
+    try:
+        progress(0, desc="Initializing...")
+        
+        # Set algorithm
+        algorithm_map = {
+            "STTN (Fast)": InpaintMode.STTN,
+            "LAMA (Quality)": InpaintMode.LAMA,
+            "ProPainter (Best)": InpaintMode.PROPAINTER,
+            "Stable Diffusion": InpaintMode.STABLE_DIFFUSION,
+            "DiffuEraser (Recommended)": InpaintMode.DIFFUERASER,
+        }
+        
+        config.MODE = algorithm_map[algorithm]
+        
+        # Apply algorithm-specific settings
+        if config.MODE == InpaintMode.STTN:
+            config.STTN_SKIP_DETECTION = sttn_skip_detection
+            config.STTN_MAX_LOAD_NUM = sttn_max_load
+        elif config.MODE == InpaintMode.LAMA:
+            config.LAMA_SUPER_FAST = lama_super_fast
+        elif config.MODE == InpaintMode.STABLE_DIFFUSION:
+            config.SD_STEPS = sd_steps
+            config.SD_GUIDANCE_SCALE = sd_guidance
+        elif config.MODE == InpaintMode.DIFFUERASER:
+            config.DIFFUERASER_STEPS = diffueraser_steps
+            config.DIFFUERASER_MAX_LOAD_NUM = diffueraser_max_load
+        
+        progress(0.1, desc="Loading video...")
+        
+        # Get video path
+        video_path = video_file if isinstance(video_file, str) else video_file.name
+        
+        # Create output filename
+        base_name = os.path.basename(video_path)
+        output_name = base_name.replace('.mp4', '_no_sub.mp4')
+        output_path = os.path.join(os.path.dirname(video_path), output_name)
+        
+        progress(0.2, desc=f"Processing with {algorithm}...")
+        
+        # Run subtitle removal
+        remover = SubtitleRemover(video_path, gui_mode=False)
+        remover.run()
+        
+        progress(1.0, desc="Complete!")
+        
+        return remover.video_out_name, f"✅ Success! Processed with {algorithm}"
+        
+    except Exception as e:
+        import traceback
+        error_msg = f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
+        return None, error_msg
+
+
+def create_interface():
+    """Create Gradio interface"""
+    
+    with gr.Blocks(title="Video Subtitle Remover", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🎬 Video Subtitle Remover
+        
+        Remove hardcoded subtitles from videos using AI inpainting.
+        
+        **⭐ Recommended**: STTN for speed (Colab T4), DiffuEraser for best quality (requires 12GB+ VRAM).
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                # Input
+                video_input = gr.Video(
+                    label="📹 Upload Video"
+                )
+                
+                algorithm = gr.Radio(
+                    choices=[
+                        "STTN (Fast)",
+                        "LAMA (Quality)",
+                        "ProPainter (Best)",
+                        "Stable Diffusion",
+                        "DiffuEraser (Recommended)"
+                    ],
+                    value="STTN (Fast)",
+                    label="🎨 Algorithm",
+                    info="STTN: Fastest. DiffuEraser: Best quality (needs 12GB+ VRAM)."
+                )
+                
+                with gr.Accordion("⚙️ Advanced Settings", open=False):
+                    with gr.Tab("STTN"):
+                        sttn_skip = gr.Checkbox(
+                            label="Skip Detection (Faster)",
+                            value=True,
+                            info="Process entire subtitle area without detection"
+                        )
+                        sttn_max_load = gr.Slider(
+                            minimum=20, maximum=100, value=40, step=10,
+                            label="Max Frames per Batch",
+                            info="Lower if OOM"
+                        )
+                    
+                    with gr.Tab("LAMA"):
+                        lama_fast = gr.Checkbox(
+                            label="Super Fast Mode",
+                            value=False,
+                            info="Use OpenCV instead of deep learning"
+                        )
+                    
+                    with gr.Tab("Stable Diffusion"):
+                        sd_steps = gr.Slider(
+                            minimum=20, maximum=100, value=50, step=5,
+                            label="Inference Steps",
+                            info="More = better quality but slower"
+                        )
+                        sd_guidance = gr.Slider(
+                            minimum=5.0, maximum=15.0, value=7.5, step=0.5,
+                            label="Guidance Scale"
+                        )
+                    
+                    with gr.Tab("DiffuEraser"):
+                        diffueraser_steps = gr.Slider(
+                            minimum=20, maximum=100, value=50, step=5,
+                            label="Diffusion Steps"
+                        )
+                        diffueraser_max_load = gr.Slider(
+                            minimum=20, maximum=120, value=40, step=10,
+                            label="Max Frames per Batch",
+                            info="T4 GPU: use 40. Higher GPU: 80+"
+                        )
+                
+                process_btn = gr.Button("🚀 Remove Subtitles", variant="primary", size="lg")
+            
+            with gr.Column():
+                # Output
+                video_output = gr.Video(label="✨ Output Video")
+                status_output = gr.Textbox(label="📊 Status", lines=3)
+                
+                gr.Markdown("""
+                ### 💡 Performance (Colab T4):
+                - **STTN**: ~30s for 1min 720p ⚡
+                - **LAMA**: ~60s for 1min 720p
+                - **DiffuEraser**: ~3-5min for 1min 720p (needs 12GB+ VRAM)
+                - **SD**: ~2-3min for 1min 720p
+                
+                ### ⚠️ Tips:
+                - Keep videos **under 10 minutes** on free Colab
+                - Use STTN for quick tests
+                - Use DiffuEraser for final high-quality output
+                - Monitor VRAM with `!nvidia-smi`
+                """)
+        
+        # Connect button
+        process_btn.click(
+            fn=process_video,
+            inputs=[
+                video_input, algorithm,
+                sttn_skip, sttn_max_load,
+                lama_fast,
+                sd_steps, sd_guidance,
+                diffueraser_steps, diffueraser_max_load
+            ],
+            outputs=[video_output, status_output]
+        )
+        
+        # Examples
+        gr.Examples(
+            examples=[
+                ["STTN (Fast)", True, 40, False, 50, 7.5, 50, 40],
+                ["DiffuEraser (Recommended)", True, 80, False, 50, 7.5, 50, 40],
+            ],
+            inputs=[
+                algorithm,
+                sttn_skip, sttn_max_load,
+                lama_fast,
+                sd_steps, sd_guidance,
+                diffueraser_steps, diffueraser_max_load
+            ],
+            label="Quick Presets"
+        )
+    
+    return demo
+
+
+if __name__ == "__main__":
+    # Check environment
+    if IN_COLAB:
+        print("🌐 Running in Google Colab")
+        print(f"📁 Project root: {PROJECT_ROOT}")
+    else:
+        print("💻 Running locally")
+    
+    # Create and launch interface
+    demo = create_interface()
+    demo.launch(
+        share=True,  # Create public URL
+        debug=True,
+        server_name="0.0.0.0",
+        server_port=7860
+    )