File size: 7,612 Bytes
0881e36 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
#!/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")
|