google_colabs/gradio_app.py

"""
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
from backend.subtitle_extractor import SubtitleExtractor


def process_video(
    video_file,
    algorithm,
    extract_subtitles,
    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 and optionally extract them"""
    if video_file is None:
        return None, 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()
        
        srt_file = None
        
        # Extract subtitles if requested
        if extract_subtitles:
            progress(0.8, desc="Extracting subtitles...")
            try:
                extractor = SubtitleExtractor(video_path)
                srt_file = extractor.extract_to_srt(
                    progress_callback=lambda p, desc: progress(0.8 + p * 0.2, desc)
                )
            except Exception as e:
                print(f"Warning: Subtitle extraction failed: {e}")
                srt_file = None
        
        progress(1.0, desc="Complete!")
        
        status_msg = f"✅ Success! Processed with {algorithm}"
        if srt_file:
            status_msg += f"\n📄 Subtitles extracted to SRT file"
        
        return remover.video_out_name, srt_file, status_msg
        
    except Exception as e:
        import traceback
        error_msg = f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
        return None, 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)."
                )
                
                extract_subs = gr.Checkbox(
                    label="📝 Extract Subtitles to SRT",
                    value=False,
                    info="Extract detected subtitles using OCR and save as SRT file"
                )
                
                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")
                srt_output = gr.File(label="📄 Extracted Subtitles (SRT)", visible=True)
                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, extract_subs,
                sttn_skip, sttn_max_load,
                lama_fast,
                sd_steps, sd_guidance,
                diffueraser_steps, diffueraser_max_load
            ],
            outputs=[video_output, srt_output, status_output]
        )
        
        # Examples
        gr.Examples(
            examples=[
                ["STTN (Fast)", False, True, 40, False, 50, 7.5, 50, 40],
                ["DiffuEraser (Recommended)", True, True, 80, False, 50, 7.5, 50, 40],
            ],
            inputs=[
                algorithm, extract_subs,
                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
    )