#!/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")