google_colabs/fix_notebook.py

#!/usr/bin/env python3
"""
Update Colab notebook with new Stable Diffusion, DiffuEraser, and E2FGVI options
"""
import json

# Read the notebook
with open('Video_Subtitle_Remover.ipynb', 'r') as f:
    nb = json.load(f)

# 1. Update Step 3 - Add new dependencies
for i, cell in enumerate(nb['cells']):
    if cell['cell_type'] == 'code' and cell['source']:
        source = ''.join(cell['source'])
        if 'pip install -q filesplit==3.0.2' in source:
            # Add diffusers and related packages
            cell['source'] = [
                "# Install dependencies\n",
                "# Note: Colab already has torch, torchvision, opencv-python, numpy, etc.\n",
                "!pip install -q filesplit==3.0.2 albumentations scikit-image imgaug pyclipper lmdb\n",
                "!pip install -q PyYAML omegaconf tqdm easydict scikit-learn pandas webdataset\n",
                "!pip install -q protobuf av einops paddleocr paddle2onnx onnxruntime-gpu\n",
                "\n",
                "# Install PaddlePaddle GPU version (compatible with Colab)\n",
                "!pip install -q paddlepaddle-gpu==2.6.2\n",
                "\n",
                "# Advanced Inpainting Models (Optional - only install if using these modes)\n",
                "# Uncomment the line below to enable Stable Diffusion, DiffuEraser, E2FGVI\n",
                "# !pip install -q diffusers transformers accelerate\n",
                "\n",
                "print(\"✓ Dependencies installed!\")"
            ]
            print(f"✓ Updated Step 3 (Cell {i}): Added advanced model dependencies")
            break

# 2. Update Step 5 - Add new algorithm options
for i, cell in enumerate(nb['cells']):
    if cell['cell_type'] == 'code' and cell['source']:
        source = ''.join(cell['source'])
        if 'ALGORITHM =' in source and '# === CONFIGURATION ===' in source:
            cell['source'] = [
                "# === CONFIGURATION ===\n",
                "\n",
                "# Algorithm selection\n",
                "# Options:\n",
                "#   'STTN' - Fast, real-time video (recommended for Colab)\n",
                "#   'LAMA' - High quality for images/animation\n",
                "#   'PROPAINTER' - Best quality, very slow, high VRAM\n",
                "#   'SD' - Stable Diffusion (NEW - requires extra install above)\n",
                "#   'DIFFUERASER' - Specialized subtitle removal (Coming soon)\n",
                "#   'E2FGVI' - Fast flow-guided (Coming soon)\n",
                "ALGORITHM = 'STTN'\n",
                "\n",
                "# STTN Settings (recommended for Colab)\n",
                "STTN_SKIP_DETECTION = True  # Much faster, processes entire subtitle area\n",
                "STTN_MAX_LOAD_NUM = 40      # Reduce if OOM (30-50 for T4 GPU)\n",
                "STTN_NEIGHBOR_STRIDE = 5\n",
                "STTN_REFERENCE_LENGTH = 10\n",
                "\n",
                "# LAMA Settings\n",
                "LAMA_SUPER_FAST = False     # Set True for faster but lower quality\n",
                "\n",
                "# ProPainter Settings (requires 16GB+ GPU, not recommended for Colab)\n",
                "PROPAINTER_MAX_LOAD_NUM = 40  # Very low for T4 GPU\n",
                "\n",
                "# Stable Diffusion Settings (NEW)\n",
                "SD_STEPS = 50                # More steps = better quality but slower\n",
                "SD_GUIDANCE_SCALE = 7.5      # How much to follow the prompt\n",
                "SD_PROMPT = \"natural scene, high quality\"  # Text guidance\n",
                "\n",
                "# Video path (change this)\n",
                "# Option 1: Use sample video\n",
                "VIDEO_PATH = '/content/video-subtitle-remover/test/test.mp4'\n",
                "\n",
                "# Option 2: Use video from Google Drive (uncomment)\n",
                "# VIDEO_PATH = '/content/drive/MyDrive/my_video.mp4'\n",
                "\n",
                "# Subtitle area (optional, in pixels: ymin, ymax, xmin, xmax)\n",
                "# None = auto-detect subtitle area\n",
                "SUBTITLE_AREA = None\n",
                "\n",
                "# Example: Bottom 20% of 1080p video\n",
                "# SUBTITLE_AREA = (864, 1080, 0, 1920)\n",
                "\n",
                "print(f\"Configuration:\")\n",
                "print(f\"  Algorithm: {ALGORITHM}\")\n",
                "print(f\"  Video: {VIDEO_PATH}\")\n",
                "print(f\"  Subtitle area: {SUBTITLE_AREA or 'Auto-detect'}\")"
            ]
            print(f"✓ Updated Step 5 (Cell {i}): Added SD/DiffuEraser/E2FGVI options")
            break

# 3. Update Step 6 - Handle new algorithms
for i, cell in enumerate(nb['cells']):
    if cell['cell_type'] == 'code' and cell['source']:
        source = ''.join(cell['source'])
        if 'config.MODE = InpaintMode.STTN' in source:
            cell['source'] = [
                "# Modify config.py with our settings\n",
                "import sys\n",
                "sys.path.insert(0, '/content/video-subtitle-remover')\n",
                "sys.path.insert(0, '/content/video-subtitle-remover/backend')\n",
                "\n",
                "from backend import config\n",
                "from backend.config import InpaintMode\n",
                "\n",
                "# Apply algorithm selection\n",
                "if ALGORITHM == 'STTN':\n",
                "    config.MODE = InpaintMode.STTN\n",
                "    config.STTN_SKIP_DETECTION = STTN_SKIP_DETECTION\n",
                "    config.STTN_MAX_LOAD_NUM = STTN_MAX_LOAD_NUM\n",
                "    config.STTN_NEIGHBOR_STRIDE = STTN_NEIGHBOR_STRIDE\n",
                "    config.STTN_REFERENCE_LENGTH = STTN_REFERENCE_LENGTH\n",
                "elif ALGORITHM == 'LAMA':\n",
                "    config.MODE = InpaintMode.LAMA\n",
                "    config.LAMA_SUPER_FAST = LAMA_SUPER_FAST\n",
                "elif ALGORITHM == 'PROPAINTER':\n",
                "    config.MODE = InpaintMode.PROPAINTER\n",
                "    config.PROPAINTER_MAX_LOAD_NUM = PROPAINTER_MAX_LOAD_NUM\n",
                "elif ALGORITHM == 'SD':\n",
                "    config.MODE = InpaintMode.STABLE_DIFFUSION\n",
                "    config.SD_STEPS = SD_STEPS\n",
                "    config.SD_GUIDANCE_SCALE = SD_GUIDANCE_SCALE\n",
                "    config.SD_PROMPT = SD_PROMPT\n",
                "elif ALGORITHM == 'DIFFUERASER':\n",
                "    config.MODE = InpaintMode.DIFFUERASER\n",
                "    print('⚠️  DiffuEraser not yet implemented, will fall back to LAMA')\n",
                "elif ALGORITHM == 'E2FGVI':\n",
                "    config.MODE = InpaintMode.E2FGVI\n",
                "    print('⚠️  E2FGVI not yet implemented, will fall back to STTN')\n",
                "\n",
                "print(f\"✓ Configuration applied!\")\n",
                "print(f\"  Using device: {config.device}\")\n",
                "print(f\"  Mode: {config.MODE.value}\")"
            ]
            print(f"✓ Updated Step 6 (Cell {i}): Added new algorithm handling")
            break

# Save the notebook
with open('Video_Subtitle_Remover.ipynb', 'w') as f:
    json.dump(nb, f, indent=1, ensure_ascii=False)

print("\n✓ Colab notebook updated!")
print("\nChanges made:")
print("  1. Step 3: Added diffusers/transformers (commented out by default)")
print("  2. Step 5: Added SD, DiffuEraser, E2FGVI options")
print("  3. Step 5: Added SD configuration parameters")
print("  4. Step 6: Added handling for new algorithms")