mirror of
https://github.com/digitalsamba/claude-code-video-toolkit.git
synced 2026-09-18 19:41:13 +08:00
Add dewatermark and locate_watermark tools (v0.7.0)
AI-powered watermark removal using ProPainter video inpainting: dewatermark.py (2400+ lines): - Cloud processing via RunPod serverless GPUs - Local processing for NVIDIA GPU users - Auto resize-ratio based on video duration (1.0 for <30s, 0.75 for <1min) - Cloudflare R2 for reliable file transfer - Automated --setup for one-command RunPod configuration - Chunked processing for long videos - Preset watermark regions (notebooklm, tiktok, stock variants) locate_watermark.py: - Coordinate grid overlay for identifying watermark positions - Region verification across multiple frames - Preset support for common watermarks Infrastructure: - docker/runpod-propainter/ - Serverless Docker image + handler - docs/runpod-setup.md - Cloud GPU setup guide - docs/optional-components.md - ML component documentation - R2 + RunPod configuration in .env.example 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Conal Mullan
@@ -6,3 +6,23 @@ ELEVENLABS_API_KEY=your_api_key_here
|
||||
# Find voice IDs at https://elevenlabs.io/app/voice-library
|
||||
# Or create a voice clone and use that ID
|
||||
ELEVENLABS_VOICE_ID=your_voice_id_here
|
||||
|
||||
# RunPod API key for cloud GPU processing (dewatermark tool)
|
||||
# Get yours at https://runpod.io/console/user/settings
|
||||
# See docs/runpod-setup.md for full setup guide
|
||||
RUNPOD_API_KEY=your_api_key_here
|
||||
|
||||
# RunPod serverless endpoint ID
|
||||
# Create endpoint at https://runpod.io/console/serverless
|
||||
# Use the Docker image from docker/runpod-propainter/
|
||||
RUNPOD_ENDPOINT_ID=your_endpoint_id_here
|
||||
|
||||
# Cloudflare R2 - Reliable file transfer for RunPod jobs (recommended)
|
||||
# Free tier: 10GB storage, 10M ops/month, zero egress fees, no expiration
|
||||
# Setup: https://dash.cloudflare.com → R2 Object Storage → Create bucket
|
||||
# Then: Manage R2 API Tokens → Create API token (Object Read & Write)
|
||||
# If not configured, falls back to free file hosting services (less reliable)
|
||||
R2_ACCOUNT_ID=your_account_id_here
|
||||
R2_ACCESS_KEY_ID=your_access_key_id_here
|
||||
R2_SECRET_ACCESS_KEY=your_secret_access_key_here
|
||||
R2_BUCKET_NAME=video-toolkit
|
||||
|
||||
@@ -89,3 +89,4 @@ examples/*/out/
|
||||
# Playwright recordings (can be large)
|
||||
playwright/output/*.mp4
|
||||
playwright/recordings/*.mp4
|
||||
.ai_dev/
|
||||
|
||||
@@ -181,10 +181,93 @@ python tools/addmusic.py --input video.mp4 --music bg.mp3 --music-volume 0.2 --f
|
||||
| Type | Tools | When to Use |
|
||||
|------|-------|-------------|
|
||||
| **Project tools** | voiceover, music, sfx | During video creation workflow |
|
||||
| **Utility tools** | redub, addmusic, notebooklm_brand | Quick transformations on existing videos |
|
||||
| **Utility tools** | redub, addmusic, notebooklm_brand, locate_watermark | Quick transformations on existing videos |
|
||||
| **Optional tools** | dewatermark | Requires additional installation (see below) |
|
||||
|
||||
Utility tools work on any video file without requiring a project structure.
|
||||
|
||||
### Watermark Removal (Optional Component)
|
||||
|
||||
The `dewatermark.py` tool uses AI inpainting (ProPainter) to remove watermarks.
|
||||
|
||||
**Two processing modes:**
|
||||
- **RunPod (cloud)** - Works from any machine, ~$0.05-0.30/video
|
||||
- **Local** - Requires NVIDIA GPU with 8GB+ VRAM
|
||||
|
||||
```bash
|
||||
# Cloud processing via RunPod (recommended for Mac users)
|
||||
# Default outputs at 50% resolution for memory safety
|
||||
python tools/dewatermark.py --input video.mp4 --region 1080,660,195,40 --output clean.mp4 --runpod
|
||||
|
||||
# Full resolution (may fail on some GPUs due to memory limits)
|
||||
python tools/dewatermark.py --input video.mp4 --region 1080,660,195,40 --output clean.mp4 --runpod --resize-ratio 1.0
|
||||
|
||||
# Local processing (requires NVIDIA GPU + ProPainter installation)
|
||||
python tools/dewatermark.py --input video.mp4 --region 1080,660,195,40 --output clean.mp4
|
||||
|
||||
# Check local installation status
|
||||
python tools/dewatermark.py --status
|
||||
|
||||
# Install ProPainter for local processing (~2GB download)
|
||||
python tools/dewatermark.py --install
|
||||
```
|
||||
|
||||
**RunPod setup (for cloud processing):**
|
||||
```bash
|
||||
# 1. Add API key to .env
|
||||
echo "RUNPOD_API_KEY=your_key_here" >> .env
|
||||
|
||||
# 2. Run automated setup
|
||||
python tools/dewatermark.py --setup
|
||||
|
||||
# Done! The endpoint ID is automatically saved to .env
|
||||
```
|
||||
|
||||
For manual setup or advanced options, see `docs/runpod-setup.md`.
|
||||
|
||||
**Hardware requirements (local mode):**
|
||||
- **Required:** NVIDIA GPU (8GB+ VRAM)
|
||||
- **Not supported:** Apple Silicon, CPU-only (use `--runpod` instead)
|
||||
- **Disk:** ~2GB for model weights
|
||||
|
||||
**Installation location:** `~/.video-toolkit/propainter/`
|
||||
|
||||
### Locating Watermarks
|
||||
|
||||
Before removing a watermark, you need to identify its exact coordinates. The `locate_watermark.py` tool helps with this.
|
||||
|
||||
**Requires:** ImageMagick (`brew install imagemagick`)
|
||||
|
||||
```bash
|
||||
# Explore with coordinate grid overlay
|
||||
python tools/locate_watermark.py --input video.mp4 --grid --output-dir ./review/
|
||||
|
||||
# Use a preset for common watermarks
|
||||
python tools/locate_watermark.py --input video.mp4 --preset notebooklm --verify
|
||||
|
||||
# Verify custom region across multiple frames
|
||||
python tools/locate_watermark.py --input video.mp4 --region 1100,650,150,50 --verify
|
||||
|
||||
# List available presets
|
||||
python tools/locate_watermark.py --list-presets
|
||||
```
|
||||
|
||||
**Workflow:**
|
||||
1. Extract frames with `--grid` to identify watermark position
|
||||
2. Note coordinates from the grid overlay
|
||||
3. Verify with `--region x,y,w,h --verify` across multiple frames
|
||||
4. Use confirmed region with `dewatermark.py`
|
||||
|
||||
**Presets:** notebooklm, tiktok, stock-br, stock-bl, stock-center
|
||||
|
||||
**Options:**
|
||||
- `--samples N` - Number of frames to extract (default: 5)
|
||||
- `--grid` - Overlay coordinate grid
|
||||
- `--mark` - Draw rectangle on frames
|
||||
- `--verify` - Mark region across multiple frames for verification
|
||||
- `--crop` - Also output cropped watermark regions
|
||||
- `--open` - Open output directory in Finder (macOS)
|
||||
|
||||
### Redub Sync Mode
|
||||
|
||||
The `--sync` flag enables word-level time remapping for redubbing. This is essential when the TTS voice speaks at a different pace than the original.
|
||||
@@ -501,3 +584,4 @@ Keep these separate. Don't mix toolkit improvements with video production.
|
||||
- `docs/getting-started.md` - First video walkthrough
|
||||
- `docs/creating-templates.md` - Build new templates
|
||||
- `docs/creating-brands.md` - Create brand profiles
|
||||
- `docs/optional-components.md` - Setup for optional ML-based tools (ProPainter, etc.)
|
||||
|
||||
@@ -162,6 +162,47 @@
|
||||
"status": "beta",
|
||||
"created": "2025-12-28",
|
||||
"updated": "2025-12-28"
|
||||
},
|
||||
"dewatermark": {
|
||||
"path": "tools/dewatermark.py",
|
||||
"description": "Remove watermarks using AI inpainting (ProPainter) - supports local and RunPod cloud",
|
||||
"usage": "python tools/dewatermark.py --input video.mp4 --region 1080,660,195,40 --output clean.mp4 --runpod",
|
||||
"status": "beta",
|
||||
"optional": true,
|
||||
"requires": "RunPod account OR local NVIDIA GPU + ProPainter (~2GB)",
|
||||
"created": "2025-12-29",
|
||||
"updated": "2025-12-30"
|
||||
}
|
||||
},
|
||||
|
||||
"optionalComponents": {
|
||||
"propainter": {
|
||||
"name": "ProPainter",
|
||||
"description": "AI video inpainting for watermark removal (local mode)",
|
||||
"installPath": "~/.video-toolkit/propainter/",
|
||||
"repository": "https://github.com/sczhou/ProPainter.git",
|
||||
"diskSpace": "~2GB",
|
||||
"hardware": {
|
||||
"recommended": "NVIDIA GPU (8GB+ VRAM)",
|
||||
"notSupported": "Apple Silicon, CPU (use --runpod instead)",
|
||||
"memory": "8-10GB VRAM for 720p with fp16"
|
||||
},
|
||||
"installCommand": "python tools/dewatermark.py --install",
|
||||
"statusCommand": "python tools/dewatermark.py --status",
|
||||
"documentation": "docs/optional-components.md"
|
||||
}
|
||||
},
|
||||
|
||||
"cloudProviders": {
|
||||
"runpod": {
|
||||
"name": "RunPod",
|
||||
"description": "Serverless GPU cloud for video processing",
|
||||
"dockerImage": "docker/runpod-propainter/",
|
||||
"documentation": "docs/runpod-setup.md",
|
||||
"operations": ["dewatermark"],
|
||||
"envVars": ["RUNPOD_API_KEY", "RUNPOD_ENDPOINT_ID"],
|
||||
"estimatedCost": "$0.05-0.30 per video",
|
||||
"created": "2025-12-30"
|
||||
}
|
||||
},
|
||||
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
# RunPod Serverless handler for ProPainter (dewatermark)
|
||||
#
|
||||
# Build: docker build -t yourusername/video-toolkit-propainter:latest .
|
||||
# Push: docker push yourusername/video-toolkit-propainter:latest
|
||||
#
|
||||
# Image size: ~4GB (includes pre-baked model weights for fast cold starts)
|
||||
#
|
||||
# Version: 2.0.0 - CUDA 12.4, PyTorch 2.4, fixed GPU detection
|
||||
|
||||
FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
|
||||
|
||||
# Prevent interactive prompts during package installation
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install Python and system dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
python3.10 \
|
||||
python3-pip \
|
||||
python3.10-venv \
|
||||
git \
|
||||
ffmpeg \
|
||||
curl \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& ln -sf /usr/bin/python3.10 /usr/bin/python3 \
|
||||
&& ln -sf /usr/bin/python3.10 /usr/bin/python
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install PyTorch with CUDA 12.4 support
|
||||
RUN pip3 install --no-cache-dir \
|
||||
torch==2.4.1 \
|
||||
torchvision==0.19.1 \
|
||||
--index-url https://download.pytorch.org/whl/cu124
|
||||
|
||||
# Clone ProPainter repository
|
||||
RUN git clone --depth 1 https://github.com/sczhou/ProPainter.git /app/propainter
|
||||
|
||||
# Install ProPainter requirements
|
||||
# PyTorch 2.4+ has better numpy 2.x compatibility, but ProPainter may still need numpy 1.x
|
||||
WORKDIR /app/propainter
|
||||
RUN pip3 install --no-cache-dir "numpy>=1.26,<2" && \
|
||||
pip3 install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Pre-download model weights (baked into image for ~30s cold starts vs ~2min)
|
||||
# Weights are ~2GB total
|
||||
RUN mkdir -p /app/propainter/weights && \
|
||||
echo "Downloading ProPainter.pth..." && \
|
||||
curl -L -o /app/propainter/weights/ProPainter.pth \
|
||||
"https://github.com/sczhou/ProPainter/releases/download/v0.1.0/ProPainter.pth" && \
|
||||
echo "Downloading recurrent_flow_completion.pth..." && \
|
||||
curl -L -o /app/propainter/weights/recurrent_flow_completion.pth \
|
||||
"https://github.com/sczhou/ProPainter/releases/download/v0.1.0/recurrent_flow_completion.pth" && \
|
||||
echo "Downloading raft-things.pth..." && \
|
||||
curl -L -o /app/propainter/weights/raft-things.pth \
|
||||
"https://github.com/sczhou/ProPainter/releases/download/v0.1.0/raft-things.pth" && \
|
||||
echo "Downloading i3d_rgb_imagenet.pt..." && \
|
||||
curl -L -o /app/propainter/weights/i3d_rgb_imagenet.pt \
|
||||
"https://github.com/sczhou/ProPainter/releases/download/v0.1.0/i3d_rgb_imagenet.pt" && \
|
||||
echo "All weights downloaded successfully"
|
||||
|
||||
# Install RunPod SDK and additional utilities
|
||||
RUN pip3 install --no-cache-dir \
|
||||
runpod>=1.7.0 \
|
||||
requests>=2.31.0 \
|
||||
boto3>=1.34.0
|
||||
|
||||
# Copy handler
|
||||
WORKDIR /app
|
||||
COPY handler.py /app/handler.py
|
||||
|
||||
# Environment
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
# NOTE: Do NOT set CUDA_VISIBLE_DEVICES here - RunPod sets this dynamically
|
||||
# to assign the correct GPU to each serverless worker
|
||||
|
||||
# Health check - verify CUDA is available
|
||||
RUN python3 -c "import torch; print(f'PyTorch {torch.__version__}, CUDA available: {torch.cuda.is_available()}')"
|
||||
|
||||
# Run handler
|
||||
CMD ["python3", "-u", "/app/handler.py"]
|
||||
@@ -0,0 +1,235 @@
|
||||
# RunPod ProPainter Docker Image
|
||||
|
||||
Serverless GPU handler for video watermark removal using ProPainter AI inpainting.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Option A: Use Pre-built Public Image (Recommended)
|
||||
|
||||
A public image is available on GitHub Container Registry:
|
||||
|
||||
```
|
||||
ghcr.io/conalmullan/video-toolkit-propainter:latest
|
||||
```
|
||||
|
||||
Skip to **Step 2: Deploy on RunPod** below.
|
||||
|
||||
### Option B: Build Your Own Image
|
||||
|
||||
```bash
|
||||
cd docker/runpod-propainter
|
||||
|
||||
# Build for linux/amd64 (required for RunPod)
|
||||
docker buildx build --platform linux/amd64 -t yourusername/video-toolkit-propainter:latest --push .
|
||||
```
|
||||
|
||||
Build takes ~15-20 minutes (downloads ~2GB of model weights).
|
||||
|
||||
### Deploy on RunPod
|
||||
|
||||
1. Go to [RunPod Serverless](https://www.runpod.io/console/serverless)
|
||||
2. Click **New Endpoint**
|
||||
3. Configure:
|
||||
- **Docker Image**: `ghcr.io/conalmullan/video-toolkit-propainter:latest`
|
||||
- **GPU**: RTX 3090 or RTX 4090 (24GB VRAM recommended)
|
||||
- **Max Workers**: 1 (scale up as needed)
|
||||
- **Idle Timeout**: 5 seconds (fast scale-down)
|
||||
- **Execution Timeout**: 3600 seconds (1 hour max)
|
||||
4. Copy the **Endpoint ID** for your `.env` file
|
||||
|
||||
### Configure Local Tool
|
||||
|
||||
Add to your `.env`:
|
||||
|
||||
```bash
|
||||
RUNPOD_API_KEY=your_api_key_here
|
||||
RUNPOD_ENDPOINT_ID=your_endpoint_id_here
|
||||
```
|
||||
|
||||
### Use It
|
||||
|
||||
```bash
|
||||
python tools/dewatermark.py \
|
||||
--input video.mp4 \
|
||||
--region 1080,660,195,40 \
|
||||
--output clean.mp4 \
|
||||
--runpod
|
||||
```
|
||||
|
||||
## Image Details
|
||||
|
||||
| Property | Value |
|
||||
|----------|-------|
|
||||
| Base | `nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04` |
|
||||
| Size | ~4GB |
|
||||
| Cold Start | ~30 seconds |
|
||||
| Python | 3.10 |
|
||||
| PyTorch | 2.1.0 + CUDA 12.1 |
|
||||
|
||||
### Pre-baked Components
|
||||
|
||||
- ProPainter repository
|
||||
- Model weights (~2GB):
|
||||
- ProPainter.pth
|
||||
- recurrent_flow_completion.pth
|
||||
- raft-things.pth
|
||||
- i3d_rgb_imagenet.pt
|
||||
- FFmpeg for video processing
|
||||
- RunPod SDK
|
||||
|
||||
## API Reference
|
||||
|
||||
### Input Format
|
||||
|
||||
```json
|
||||
{
|
||||
"input": {
|
||||
"operation": "dewatermark",
|
||||
"video_url": "https://example.com/video.mp4",
|
||||
"region": "1080,660,195,40"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Or with a pre-made mask:
|
||||
|
||||
```json
|
||||
{
|
||||
"input": {
|
||||
"operation": "dewatermark",
|
||||
"video_url": "https://example.com/video.mp4",
|
||||
"mask_url": "https://example.com/mask.png"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Output Format
|
||||
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"output_url": "https://runpod-storage.../job123_dewatermarked.mp4",
|
||||
"video_dimensions": "1920x1080",
|
||||
"video_duration_seconds": 45.5,
|
||||
"gpu_vram_gb": 24,
|
||||
"profile_used": {
|
||||
"subvideo_length": 60,
|
||||
"neighbor_length": 10,
|
||||
"ref_stride": 10
|
||||
},
|
||||
"processing_time_seconds": 120.5
|
||||
}
|
||||
```
|
||||
|
||||
### Error Format
|
||||
|
||||
```json
|
||||
{
|
||||
"error": "Description of what went wrong"
|
||||
}
|
||||
```
|
||||
|
||||
## GPU Memory Profiles
|
||||
|
||||
The handler auto-detects GPU VRAM and selects optimal settings:
|
||||
|
||||
| VRAM | subvideo_length | neighbor_length | ref_stride | Speed |
|
||||
|------|-----------------|-----------------|------------|-------|
|
||||
| 8GB | 30 | 5 | 25 | Slow |
|
||||
| 12GB | 40 | 5 | 20 | Medium |
|
||||
| 16GB | 50 | 8 | 15 | Good |
|
||||
| 24GB | 60 | 10 | 10 | Fast |
|
||||
| 48GB | 80 | 10 | 10 | Fastest |
|
||||
|
||||
**Recommendation**: Use RTX 3090 or RTX 4090 (24GB) for best price/performance.
|
||||
|
||||
## Cost Estimates
|
||||
|
||||
Using RTX 3090 (~$0.34/hr):
|
||||
|
||||
| Video Length | Processing Time | Estimated Cost |
|
||||
|--------------|-----------------|----------------|
|
||||
| < 30 seconds | 2-5 minutes | ~$0.02 |
|
||||
| 30s - 2 min | 5-15 minutes | ~$0.08 |
|
||||
| 2 - 5 min | 15-45 minutes | ~$0.25 |
|
||||
| > 5 min | 45+ minutes | ~$0.40+ |
|
||||
|
||||
## Local Testing
|
||||
|
||||
Test the image locally with NVIDIA GPU:
|
||||
|
||||
```bash
|
||||
# Build
|
||||
docker build -t propainter-test .
|
||||
|
||||
# Run interactive shell
|
||||
docker run --gpus all -it propainter-test /bin/bash
|
||||
|
||||
# Inside container, test GPU
|
||||
python3 -c "import torch; print(f'CUDA: {torch.cuda.is_available()}')"
|
||||
|
||||
# Test handler with mock job
|
||||
python3 -c "
|
||||
from handler import handler
|
||||
result = handler({
|
||||
'id': 'test123',
|
||||
'input': {
|
||||
'operation': 'dewatermark',
|
||||
'video_url': 'https://example.com/test.mp4',
|
||||
'region': '100,100,50,50'
|
||||
}
|
||||
})
|
||||
print(result)
|
||||
"
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### "CUDA out of memory"
|
||||
|
||||
The video is too long for the GPU. Options:
|
||||
1. Use a GPU with more VRAM
|
||||
2. The local tool will auto-chunk long videos (coming soon)
|
||||
|
||||
### "Failed to download video"
|
||||
|
||||
- Check the video URL is publicly accessible
|
||||
- URLs must be direct downloads (not web pages)
|
||||
- Timeout is 5 minutes for download
|
||||
|
||||
### "No output file found"
|
||||
|
||||
ProPainter failed silently. Check:
|
||||
- Video format is supported (MP4, MOV, AVI)
|
||||
- Mask dimensions match video dimensions
|
||||
- Region coordinates are valid
|
||||
|
||||
### Cold start is slow
|
||||
|
||||
First request after idle takes ~30s to load models. Subsequent requests are faster. Consider:
|
||||
- Setting longer idle timeout (but costs more)
|
||||
- Using "always on" worker for frequent usage
|
||||
|
||||
## Extending for New Operations
|
||||
|
||||
The handler is designed for extensibility. To add a new GPU operation:
|
||||
|
||||
1. Add handler function in `handler.py`:
|
||||
```python
|
||||
def handle_upscale(job_input: dict, job_id: str, work_dir: Path) -> dict:
|
||||
# Implementation
|
||||
pass
|
||||
```
|
||||
|
||||
2. Register in main handler:
|
||||
```python
|
||||
if operation == "upscale":
|
||||
return handle_upscale(job_input, job_id, work_dir)
|
||||
```
|
||||
|
||||
3. Rebuild and push image
|
||||
|
||||
Future operations might include:
|
||||
- `upscale` - Video upscaling with Real-ESRGAN
|
||||
- `denoise` - Audio/video denoising
|
||||
- `stabilize` - Video stabilization
|
||||
@@ -0,0 +1,627 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
RunPod serverless handler for video toolkit GPU operations.
|
||||
|
||||
Currently supports:
|
||||
- dewatermark: Remove watermarks using ProPainter AI inpainting
|
||||
|
||||
Extensible design for future GPU tools (upscaling, denoising, etc.).
|
||||
|
||||
Input format:
|
||||
{
|
||||
"operation": "dewatermark",
|
||||
"video_url": "https://...",
|
||||
"region": "x,y,width,height", # OR
|
||||
"mask_url": "https://..." # Pre-made mask image
|
||||
}
|
||||
|
||||
Output format:
|
||||
{
|
||||
"success": true,
|
||||
"output_url": "https://...",
|
||||
"video_dimensions": "1920x1080",
|
||||
"gpu_vram_gb": 24,
|
||||
"processing_time_seconds": 120.5
|
||||
}
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
import runpod
|
||||
|
||||
# ProPainter installation path (baked into Docker image)
|
||||
PROPAINTER_PATH = Path("/app/propainter")
|
||||
|
||||
# Memory profiles based on GPU VRAM (GB)
|
||||
# CONSERVATIVE settings to avoid OOM - ProPainter's RAFT optical flow is extremely memory hungry
|
||||
# Key parameters:
|
||||
# - subvideo_length: frames per batch (lower = less memory, more batches)
|
||||
# - neighbor_length: local temporal context (lower = less memory)
|
||||
# - ref_stride: global reference sampling (higher = fewer refs = less memory)
|
||||
#
|
||||
# NOTE: These are deliberately conservative. ProPainter defaults are:
|
||||
# subvideo_length=80, neighbor_length=10, ref_stride=10
|
||||
# But those OOM on long videos even with 80GB VRAM.
|
||||
MEMORY_PROFILES = {
|
||||
8: {"subvideo_length": 20, "neighbor_length": 3, "ref_stride": 30}, # 8GB - minimal
|
||||
11: {"subvideo_length": 25, "neighbor_length": 4, "ref_stride": 25}, # 12GB cards
|
||||
15: {"subvideo_length": 30, "neighbor_length": 5, "ref_stride": 20}, # 16GB cards
|
||||
22: {"subvideo_length": 40, "neighbor_length": 5, "ref_stride": 15}, # 24GB cards (3090, 4090)
|
||||
45: {"subvideo_length": 50, "neighbor_length": 5, "ref_stride": 15}, # 48GB cards (A6000, A40)
|
||||
75: {"subvideo_length": 60, "neighbor_length": 5, "ref_stride": 15}, # 80GB cards (A100, H100)
|
||||
}
|
||||
|
||||
|
||||
def log(message: str) -> None:
|
||||
"""Log message to stderr (visible in RunPod logs)."""
|
||||
print(message, file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
def get_gpu_vram_gb() -> int:
|
||||
"""Detect GPU VRAM using PyTorch (respects CUDA_VISIBLE_DEVICES set by RunPod)."""
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
# Get the current device (respects CUDA_VISIBLE_DEVICES)
|
||||
device_id = torch.cuda.current_device()
|
||||
props = torch.cuda.get_device_properties(device_id)
|
||||
vram_bytes = props.total_memory
|
||||
vram_gb = vram_bytes // (1024 ** 3)
|
||||
log(f"Detected GPU: {props.name}, VRAM: {vram_gb}GB ({vram_bytes // (1024**2)}MB)")
|
||||
return vram_gb
|
||||
else:
|
||||
log("Warning: CUDA not available")
|
||||
except Exception as e:
|
||||
log(f"Warning: Could not detect GPU VRAM via torch: {e}")
|
||||
|
||||
# Fallback to nvidia-smi if torch detection fails
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10
|
||||
)
|
||||
if result.returncode == 0:
|
||||
lines = result.stdout.strip().split('\n')
|
||||
vram_values = [int(line.strip()) for line in lines if line.strip()]
|
||||
if vram_values:
|
||||
max_vram_mb = max(vram_values)
|
||||
log(f"Fallback nvidia-smi detection: {max_vram_mb // 1024}GB")
|
||||
return max_vram_mb // 1024
|
||||
except Exception as e:
|
||||
log(f"Warning: nvidia-smi fallback also failed: {e}")
|
||||
|
||||
return 16 # Default assumption
|
||||
|
||||
|
||||
def get_memory_profile(vram_gb: int) -> dict:
|
||||
"""Get optimal ProPainter settings based on available VRAM."""
|
||||
for threshold in sorted(MEMORY_PROFILES.keys(), reverse=True):
|
||||
if vram_gb >= threshold:
|
||||
return MEMORY_PROFILES[threshold].copy()
|
||||
return MEMORY_PROFILES[8].copy()
|
||||
|
||||
|
||||
def download_file(url: str, output_path: str, description: str = "file") -> bool:
|
||||
"""Download file from URL with progress logging."""
|
||||
try:
|
||||
log(f"Downloading {description} from {url[:80]}...")
|
||||
response = requests.get(url, stream=True, timeout=300)
|
||||
response.raise_for_status()
|
||||
|
||||
total_size = int(response.headers.get('content-length', 0))
|
||||
downloaded = 0
|
||||
|
||||
with open(output_path, 'wb') as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
downloaded += len(chunk)
|
||||
if total_size > 0 and downloaded % (1024 * 1024) == 0:
|
||||
pct = (downloaded / total_size) * 100
|
||||
log(f" Downloaded {downloaded // (1024*1024)}MB ({pct:.0f}%)")
|
||||
|
||||
log(f" Downloaded {description}: {Path(output_path).stat().st_size // (1024*1024)}MB")
|
||||
return True
|
||||
except Exception as e:
|
||||
log(f"Error downloading {description}: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def upload_to_r2(file_path: str, job_id: str, r2_config: dict) -> tuple[Optional[str], Optional[str]]:
|
||||
"""Upload file to Cloudflare R2 and return (presigned_url, object_key)."""
|
||||
try:
|
||||
import boto3
|
||||
from botocore.config import Config
|
||||
import uuid
|
||||
|
||||
log(f"Uploading result to R2 ({Path(file_path).stat().st_size // (1024*1024)}MB)...")
|
||||
|
||||
client = boto3.client(
|
||||
"s3",
|
||||
endpoint_url=r2_config["endpoint_url"],
|
||||
aws_access_key_id=r2_config["access_key_id"],
|
||||
aws_secret_access_key=r2_config["secret_access_key"],
|
||||
config=Config(signature_version="s3v4"),
|
||||
)
|
||||
|
||||
object_key = f"dewatermark/results/{job_id}_{uuid.uuid4().hex[:8]}.mp4"
|
||||
|
||||
client.upload_file(file_path, r2_config["bucket_name"], object_key)
|
||||
|
||||
# Generate presigned URL (valid for 2 hours)
|
||||
presigned_url = client.generate_presigned_url(
|
||||
"get_object",
|
||||
Params={"Bucket": r2_config["bucket_name"], "Key": object_key},
|
||||
ExpiresIn=7200,
|
||||
)
|
||||
|
||||
log(f" R2 upload complete: {object_key}")
|
||||
return presigned_url, object_key
|
||||
except ImportError:
|
||||
log("Error: boto3 not available for R2 upload")
|
||||
return None, None
|
||||
except Exception as e:
|
||||
log(f"Error uploading to R2: {e}")
|
||||
return None, None
|
||||
|
||||
|
||||
def upload_file(file_path: str, job_id: str, r2_config: Optional[dict] = None) -> dict:
|
||||
"""
|
||||
Upload file and return upload info.
|
||||
|
||||
Returns dict with:
|
||||
- output_url: Presigned URL for download (always present if successful)
|
||||
- r2_key: R2 object key (only if R2 was used)
|
||||
"""
|
||||
# Try R2 first if configured
|
||||
if r2_config:
|
||||
url, r2_key = upload_to_r2(file_path, job_id, r2_config)
|
||||
if url:
|
||||
return {"output_url": url, "r2_key": r2_key}
|
||||
log("R2 upload failed, falling back to RunPod storage")
|
||||
|
||||
# Fall back to RunPod storage
|
||||
try:
|
||||
log(f"Uploading result to RunPod storage ({Path(file_path).stat().st_size // (1024*1024)}MB)...")
|
||||
|
||||
result_url = runpod.serverless.utils.rp_upload.upload_file_to_bucket(
|
||||
file_name=f"{job_id}_dewatermarked.mp4",
|
||||
file_location=file_path
|
||||
)
|
||||
|
||||
log(f" Upload complete: {result_url[:80]}...")
|
||||
return {"output_url": result_url}
|
||||
except Exception as e:
|
||||
log(f"Error uploading file: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def get_video_info(video_path: str) -> dict:
|
||||
"""Get video dimensions, duration, fps, and frame count using ffprobe."""
|
||||
try:
|
||||
result = subprocess.run([
|
||||
"ffprobe", "-v", "error",
|
||||
"-select_streams", "v:0",
|
||||
"-show_entries", "stream=width,height,duration,r_frame_rate,nb_frames",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "json",
|
||||
video_path,
|
||||
], capture_output=True, text=True, timeout=30)
|
||||
|
||||
if result.returncode == 0:
|
||||
import json
|
||||
data = json.loads(result.stdout)
|
||||
stream = data.get("streams", [{}])[0]
|
||||
fmt = data.get("format", {})
|
||||
|
||||
# Parse frame rate (can be "30/1" or "29.97")
|
||||
fps_str = stream.get("r_frame_rate", "30/1")
|
||||
if "/" in str(fps_str):
|
||||
num, den = str(fps_str).split("/")
|
||||
fps = float(num) / float(den) if float(den) != 0 else 30.0
|
||||
else:
|
||||
fps = float(fps_str) if fps_str else 30.0
|
||||
|
||||
duration = float(stream.get("duration") or fmt.get("duration", 0))
|
||||
|
||||
# Calculate frame count
|
||||
nb_frames = stream.get("nb_frames")
|
||||
if nb_frames:
|
||||
frame_count = int(nb_frames)
|
||||
else:
|
||||
frame_count = int(duration * fps)
|
||||
|
||||
return {
|
||||
"width": int(stream.get("width", 0)),
|
||||
"height": int(stream.get("height", 0)),
|
||||
"duration": duration,
|
||||
"fps": fps,
|
||||
"frame_count": frame_count,
|
||||
}
|
||||
except Exception as e:
|
||||
log(f"Warning: Could not get video info: {e}")
|
||||
return {"width": 0, "height": 0, "duration": 0, "fps": 30.0, "frame_count": 0}
|
||||
|
||||
|
||||
# Memory estimation constants (empirically determined from ProPainter testing)
|
||||
# ProPainter uses ~6.5MB per frame at 720p for RGB tensors, flow, masks, etc.
|
||||
BYTES_PER_FRAME_720P = 6.5 * 1024 * 1024
|
||||
|
||||
|
||||
def calculate_safe_resize_ratio(
|
||||
vram_gb: int,
|
||||
width: int,
|
||||
height: int,
|
||||
frame_count: int,
|
||||
requested_ratio: float = 1.0,
|
||||
safety_margin: float = 0.7,
|
||||
) -> tuple[float, str]:
|
||||
"""
|
||||
Calculate a safe resize_ratio based on available VRAM and video properties.
|
||||
|
||||
Returns (resize_ratio, reason) tuple.
|
||||
|
||||
The ratio is the MINIMUM of:
|
||||
- requested_ratio (what the user asked for)
|
||||
- calculated safe ratio (based on memory estimation)
|
||||
|
||||
Args:
|
||||
vram_gb: Available GPU VRAM in GB
|
||||
width: Video width in pixels
|
||||
height: Video height in pixels
|
||||
frame_count: Number of frames in video
|
||||
requested_ratio: User's requested resize ratio (default 1.0 = full res)
|
||||
safety_margin: Use this fraction of VRAM (default 0.7 = 70%)
|
||||
|
||||
Returns:
|
||||
(resize_ratio, reason): The ratio to use and why
|
||||
"""
|
||||
# Calculate memory needed at full resolution
|
||||
pixels = width * height
|
||||
pixels_720p = 1280 * 720
|
||||
scale_factor = pixels / pixels_720p
|
||||
|
||||
bytes_per_frame = BYTES_PER_FRAME_720P * scale_factor
|
||||
total_bytes_full_res = bytes_per_frame * frame_count
|
||||
|
||||
# Add overhead for model weights, intermediate tensors, etc. (~2GB base)
|
||||
model_overhead_bytes = 2 * (1024 ** 3)
|
||||
total_needed_full_res = total_bytes_full_res + model_overhead_bytes
|
||||
|
||||
# Available memory with safety margin
|
||||
available_bytes = vram_gb * (1024 ** 3) * safety_margin
|
||||
|
||||
# If full resolution fits, use requested ratio
|
||||
if total_needed_full_res <= available_bytes:
|
||||
log(f"Memory estimate: {total_needed_full_res / (1024**3):.1f}GB needed, {available_bytes / (1024**3):.1f}GB available - full resolution OK")
|
||||
return (requested_ratio, "full_resolution_fits")
|
||||
|
||||
# Calculate the resize ratio needed to fit in memory
|
||||
# Memory scales with resize_ratio^2 (both width and height reduced)
|
||||
# So: needed_memory * ratio^2 + overhead = available
|
||||
# ratio^2 = (available - overhead) / frame_memory
|
||||
# ratio = sqrt((available - overhead) / frame_memory)
|
||||
|
||||
frame_memory = total_bytes_full_res - model_overhead_bytes
|
||||
if frame_memory <= 0:
|
||||
return (0.5, "fallback_conservative")
|
||||
|
||||
usable_for_frames = available_bytes - model_overhead_bytes
|
||||
if usable_for_frames <= 0:
|
||||
return (0.25, "very_low_vram")
|
||||
|
||||
import math
|
||||
safe_ratio = math.sqrt(usable_for_frames / frame_memory)
|
||||
|
||||
# Clamp to reasonable range [0.25, 1.0]
|
||||
safe_ratio = max(0.25, min(1.0, safe_ratio))
|
||||
|
||||
# Use the more conservative of user request and calculated safe ratio
|
||||
final_ratio = min(requested_ratio, safe_ratio)
|
||||
|
||||
# Round to nice values for consistency
|
||||
nice_ratios = [1.0, 0.75, 0.5, 0.375, 0.25]
|
||||
for nice in nice_ratios:
|
||||
if final_ratio >= nice:
|
||||
final_ratio = nice
|
||||
break
|
||||
|
||||
log(f"Memory estimate: {total_needed_full_res / (1024**3):.1f}GB needed at full res, {available_bytes / (1024**3):.1f}GB available")
|
||||
log(f"Calculated safe ratio: {safe_ratio:.2f}, using: {final_ratio}")
|
||||
|
||||
reason = "auto_calculated" if final_ratio < requested_ratio else "user_requested"
|
||||
return (final_ratio, reason)
|
||||
|
||||
|
||||
def create_mask_from_region(region: str, width: int, height: int, output_path: str) -> bool:
|
||||
"""Create white-on-black mask image from x,y,w,h region string."""
|
||||
try:
|
||||
parts = [v.strip() for v in region.split(",")]
|
||||
if len(parts) != 4:
|
||||
log(f"Error: Region must be x,y,width,height - got: {region}")
|
||||
return False
|
||||
|
||||
x, y, w, h = [int(p) for p in parts]
|
||||
|
||||
# Validate bounds
|
||||
if x < 0 or y < 0 or w <= 0 or h <= 0:
|
||||
log(f"Error: Invalid region values: {region}")
|
||||
return False
|
||||
if x + w > width or y + h > height:
|
||||
log(f"Error: Region {region} exceeds video dimensions {width}x{height}")
|
||||
return False
|
||||
|
||||
log(f"Creating mask: {w}x{h} region at ({x},{y}) on {width}x{height} canvas")
|
||||
|
||||
result = subprocess.run([
|
||||
"ffmpeg", "-y",
|
||||
"-f", "lavfi",
|
||||
"-i", f"color=black:s={width}x{height}:d=1",
|
||||
"-vf", f"drawbox=x={x}:y={y}:w={w}:h={h}:c=white:t=fill",
|
||||
"-frames:v", "1",
|
||||
output_path,
|
||||
], capture_output=True, text=True, timeout=30)
|
||||
|
||||
if result.returncode != 0:
|
||||
log(f"FFmpeg error: {result.stderr}")
|
||||
return False
|
||||
|
||||
return Path(output_path).exists()
|
||||
except Exception as e:
|
||||
log(f"Error creating mask: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def run_propainter(
|
||||
video_path: str,
|
||||
mask_path: str,
|
||||
output_dir: str,
|
||||
profile: dict,
|
||||
fp16: bool = True,
|
||||
resize_ratio: float = 1.0
|
||||
) -> Optional[str]:
|
||||
"""Run ProPainter inference and return path to output video."""
|
||||
|
||||
inference_script = PROPAINTER_PATH / "inference_propainter.py"
|
||||
|
||||
cmd = [
|
||||
"python3", str(inference_script),
|
||||
"-i", video_path,
|
||||
"-m", mask_path,
|
||||
"-o", output_dir,
|
||||
"--neighbor_length", str(profile["neighbor_length"]),
|
||||
"--ref_stride", str(profile["ref_stride"]),
|
||||
"--subvideo_length", str(profile["subvideo_length"]),
|
||||
]
|
||||
|
||||
if fp16:
|
||||
cmd.append("--fp16")
|
||||
|
||||
if resize_ratio != 1.0:
|
||||
cmd.extend(["--resize_ratio", str(resize_ratio)])
|
||||
|
||||
log(f"Running ProPainter with settings: {profile}, resize_ratio={resize_ratio}")
|
||||
log(f"Command: {' '.join(cmd)}")
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
cwd=PROPAINTER_PATH,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=3600 # 1 hour max
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
log(f"ProPainter completed in {elapsed:.1f}s")
|
||||
|
||||
if result.returncode != 0:
|
||||
log(f"ProPainter error (exit {result.returncode}):")
|
||||
log("=== STDOUT (last 3000 chars) ===")
|
||||
log(result.stdout[-3000:] if result.stdout else "No stdout")
|
||||
log("=== STDERR (last 3000 chars) ===")
|
||||
log(result.stderr[-3000:] if result.stderr else "No stderr")
|
||||
return None
|
||||
|
||||
# Log success output for debugging
|
||||
log("=== ProPainter completed successfully ===")
|
||||
if result.stdout:
|
||||
log(f"STDOUT (last 1000 chars): {result.stdout[-1000:]}")
|
||||
|
||||
# Find output file - ProPainter creates: output_dir/video_name/inpaint_out.mp4
|
||||
video_stem = Path(video_path).stem
|
||||
expected = Path(output_dir) / video_stem / "inpaint_out.mp4"
|
||||
|
||||
if expected.exists():
|
||||
log(f"Output found: {expected}")
|
||||
return str(expected)
|
||||
|
||||
# Fallback: search for inpaint_out.mp4 anywhere in results
|
||||
for mp4 in Path(output_dir).rglob("inpaint_out.mp4"):
|
||||
log(f"Output found (search): {mp4}")
|
||||
return str(mp4)
|
||||
|
||||
# Last resort: any mp4 that's NOT masked_in.mp4
|
||||
for mp4 in Path(output_dir).rglob("*.mp4"):
|
||||
if "masked" not in mp4.name:
|
||||
log(f"Output found (fallback): {mp4}")
|
||||
return str(mp4)
|
||||
|
||||
log(f"Error: No output file found in {output_dir}")
|
||||
return None
|
||||
|
||||
|
||||
def handle_dewatermark(job_input: dict, job_id: str, work_dir: Path) -> dict:
|
||||
"""
|
||||
Handle dewatermark operation using ProPainter.
|
||||
|
||||
Required inputs:
|
||||
video_url: URL to video file
|
||||
region: "x,y,width,height" OR mask_url: URL to mask image
|
||||
|
||||
Optional inputs:
|
||||
fp16: Use half precision (default: true, faster)
|
||||
resize_ratio: Scale factor for processing (default: "auto" - calculated based on VRAM)
|
||||
Set to a specific value (0.25-1.0) to override auto-calculation
|
||||
r2: R2 config for result upload (endpoint_url, access_key_id, secret_access_key, bucket_name)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Validate inputs
|
||||
video_url = job_input.get("video_url")
|
||||
region = job_input.get("region")
|
||||
mask_url = job_input.get("mask_url")
|
||||
fp16 = job_input.get("fp16", True)
|
||||
requested_resize_ratio = job_input.get("resize_ratio", "auto") # Default to auto-calculation
|
||||
r2_config = job_input.get("r2") # Optional R2 config for result upload
|
||||
|
||||
if not video_url:
|
||||
return {"error": "Missing required 'video_url' in input"}
|
||||
if not region and not mask_url:
|
||||
return {"error": "Either 'region' (x,y,w,h) or 'mask_url' is required"}
|
||||
|
||||
if r2_config:
|
||||
log("R2 config provided - will upload result to R2")
|
||||
log(f"Processing options: fp16={fp16}, requested_resize_ratio={requested_resize_ratio}")
|
||||
|
||||
# Download video
|
||||
video_path = str(work_dir / "input_video.mp4")
|
||||
if not download_file(video_url, video_path, "video"):
|
||||
return {"error": "Failed to download video from URL"}
|
||||
|
||||
# Get video info
|
||||
video_info = get_video_info(video_path)
|
||||
width, height = video_info["width"], video_info["height"]
|
||||
duration = video_info["duration"]
|
||||
frame_count = video_info["frame_count"]
|
||||
|
||||
if not width or not height:
|
||||
return {"error": "Could not read video dimensions"}
|
||||
|
||||
log(f"Video: {width}x{height}, {duration:.1f}s, {frame_count} frames")
|
||||
|
||||
# Detect GPU and get optimal settings
|
||||
vram_gb = get_gpu_vram_gb()
|
||||
profile = get_memory_profile(vram_gb)
|
||||
log(f"GPU VRAM: {vram_gb}GB, using profile: {profile}")
|
||||
|
||||
# Calculate safe resize_ratio based on VRAM and video properties
|
||||
if requested_resize_ratio == "auto":
|
||||
# Auto mode: calculate optimal ratio, aim for full resolution if possible
|
||||
resize_ratio, resize_reason = calculate_safe_resize_ratio(
|
||||
vram_gb, width, height, frame_count, requested_ratio=1.0
|
||||
)
|
||||
else:
|
||||
# User specified a ratio - use it but warn if it might OOM
|
||||
user_ratio = float(requested_resize_ratio)
|
||||
safe_ratio, _ = calculate_safe_resize_ratio(
|
||||
vram_gb, width, height, frame_count, requested_ratio=user_ratio
|
||||
)
|
||||
if safe_ratio < user_ratio:
|
||||
log(f"WARNING: Requested resize_ratio={user_ratio} may cause OOM. Safe ratio is {safe_ratio}")
|
||||
resize_ratio = user_ratio
|
||||
resize_reason = "user_specified"
|
||||
|
||||
log(f"Using resize_ratio={resize_ratio} ({resize_reason})")
|
||||
|
||||
# Prepare mask
|
||||
mask_path = str(work_dir / "mask.png")
|
||||
|
||||
if mask_url:
|
||||
if not download_file(mask_url, mask_path, "mask"):
|
||||
return {"error": "Failed to download mask from URL"}
|
||||
else:
|
||||
if not create_mask_from_region(region, width, height, mask_path):
|
||||
return {"error": f"Failed to create mask from region: {region}"}
|
||||
|
||||
# Run ProPainter
|
||||
output_dir = str(work_dir / "results")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
result_path = run_propainter(video_path, mask_path, output_dir, profile, fp16, resize_ratio)
|
||||
|
||||
if not result_path:
|
||||
return {"error": "ProPainter processing failed - check logs for details"}
|
||||
|
||||
# Upload result (to R2 if configured, otherwise RunPod storage)
|
||||
upload_result = upload_file(result_path, job_id, r2_config)
|
||||
|
||||
if not upload_result.get("output_url"):
|
||||
return {"error": "Failed to upload result video"}
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
result = {
|
||||
"success": True,
|
||||
"output_url": upload_result["output_url"],
|
||||
"video_dimensions": f"{width}x{height}",
|
||||
"video_duration_seconds": round(duration, 2),
|
||||
"video_frame_count": frame_count,
|
||||
"gpu_vram_gb": vram_gb,
|
||||
"profile_used": profile,
|
||||
"resize_ratio": resize_ratio,
|
||||
"resize_reason": resize_reason,
|
||||
"processing_time_seconds": round(elapsed, 2),
|
||||
}
|
||||
|
||||
# Include R2 key if result was uploaded to R2
|
||||
if upload_result.get("r2_key"):
|
||||
result["r2_key"] = upload_result["r2_key"]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def handler(job: dict) -> dict:
|
||||
"""
|
||||
Main RunPod handler - routes to specific operations.
|
||||
|
||||
Supports operations:
|
||||
- dewatermark: Remove watermarks using ProPainter
|
||||
- (future: upscale, denoise, etc.)
|
||||
"""
|
||||
job_id = job.get("id", "unknown")
|
||||
job_input = job.get("input", {})
|
||||
|
||||
operation = job_input.get("operation", "dewatermark")
|
||||
log(f"Job {job_id}: operation={operation}")
|
||||
|
||||
# Create temp working directory
|
||||
work_dir = Path(tempfile.mkdtemp(prefix=f"runpod_{job_id}_"))
|
||||
log(f"Working directory: {work_dir}")
|
||||
|
||||
try:
|
||||
if operation == "dewatermark":
|
||||
return handle_dewatermark(job_input, job_id, work_dir)
|
||||
else:
|
||||
return {"error": f"Unknown operation: {operation}. Supported: dewatermark"}
|
||||
except Exception as e:
|
||||
import traceback
|
||||
log(f"Handler exception: {e}")
|
||||
log(traceback.format_exc())
|
||||
return {"error": f"Internal error: {str(e)}"}
|
||||
finally:
|
||||
# Cleanup temp files
|
||||
try:
|
||||
shutil.rmtree(work_dir, ignore_errors=True)
|
||||
log("Cleaned up working directory")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# RunPod serverless entry point
|
||||
if __name__ == "__main__":
|
||||
log("Starting RunPod ProPainter handler...")
|
||||
log(f"ProPainter path: {PROPAINTER_PATH}")
|
||||
log(f"Weights exist: {(PROPAINTER_PATH / 'weights').exists()}")
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,110 @@
|
||||
# Optional Components
|
||||
|
||||
Some toolkit features require additional software that isn't included by default. These **optional components** are:
|
||||
|
||||
- Installed on-demand (not part of base toolkit)
|
||||
- Stored in `~/.video-toolkit/` (outside the project)
|
||||
- Only needed for specific use cases
|
||||
|
||||
## Available Optional Components
|
||||
|
||||
| Component | Tool | Purpose | Size |
|
||||
|-----------|------|---------|------|
|
||||
| ProPainter | `dewatermark.py` | AI video inpainting for watermark removal | ~2GB |
|
||||
|
||||
## ProPainter (Watermark Removal)
|
||||
|
||||
[ProPainter](https://github.com/sczhou/ProPainter) is an AI video inpainting model that can intelligently remove watermarks by reconstructing the underlying content.
|
||||
|
||||
### Hardware Requirements
|
||||
|
||||
| Hardware | Status | Notes |
|
||||
|----------|--------|-------|
|
||||
| NVIDIA GPU (8GB+ VRAM) | **Supported** | Recommended, ~5-15 min per minute of video |
|
||||
| Cloud GPU (RunPod, etc.) | **Supported** | Good alternative, ~$0.20-0.50 per video |
|
||||
| Apple Silicon (M1/M2/M3/M4) | **Not supported** | MPS is too slow (40+ hours for short videos) |
|
||||
| CPU only | **Not supported** | Impractical processing times |
|
||||
|
||||
### Why Apple Silicon Doesn't Work
|
||||
|
||||
ProPainter relies on optical flow (RAFT) which performs extremely poorly on Apple's MPS backend:
|
||||
|
||||
1. **MPS INT_MAX Limit**: MPS cannot handle tensors > 2^31 elements, limiting chunks to ~32 seconds at 720p
|
||||
2. **MPS Performance**: Optical flow on MPS is orders of magnitude slower than CUDA
|
||||
3. **Real-world result**: 5 seconds of video takes 4+ hours on M1/M2/M3/M4
|
||||
|
||||
This is a PyTorch/MPS limitation, not something we can fix in the tool.
|
||||
|
||||
### Installation
|
||||
|
||||
```bash
|
||||
# Check current status
|
||||
python tools/dewatermark.py --status
|
||||
|
||||
# Install ProPainter
|
||||
python tools/dewatermark.py --install
|
||||
```
|
||||
|
||||
This will:
|
||||
1. Clone ProPainter to `~/.video-toolkit/propainter/`
|
||||
2. Create a Python virtual environment
|
||||
3. Install PyTorch and dependencies
|
||||
4. Download model weights (~2GB)
|
||||
|
||||
### Usage
|
||||
|
||||
**Remove watermark by specifying region:**
|
||||
```bash
|
||||
python tools/dewatermark.py \
|
||||
--input video.mp4 \
|
||||
--region 1080,660,195,40 \
|
||||
--output clean.mp4
|
||||
```
|
||||
|
||||
**Use a custom mask image:**
|
||||
```bash
|
||||
python tools/dewatermark.py \
|
||||
--input video.mp4 \
|
||||
--mask mask.png \
|
||||
--output clean.mp4
|
||||
```
|
||||
|
||||
### Finding Watermark Coordinates
|
||||
|
||||
Use the `locate_watermark.py` helper:
|
||||
|
||||
```bash
|
||||
# Extract frames with coordinate grid
|
||||
python tools/locate_watermark.py --input video.mp4 --grid --output-dir ./review/
|
||||
|
||||
# Verify a region across multiple frames
|
||||
python tools/locate_watermark.py --input video.mp4 --region 1100,650,150,50 --verify
|
||||
```
|
||||
|
||||
### Cloud GPU Alternative
|
||||
|
||||
For users without NVIDIA GPUs, cloud services offer affordable processing:
|
||||
|
||||
| Provider | GPU | Cost | Processing Time |
|
||||
|----------|-----|------|-----------------|
|
||||
| RunPod | RTX 4090 | ~$0.34/hr | ~15-30 min for 3-min video |
|
||||
| RunPod | A100 | ~$1.99/hr | ~5-15 min for 3-min video |
|
||||
| Vast.ai | RTX 3090 | ~$0.20/hr | ~20-40 min for 3-min video |
|
||||
|
||||
Both RunPod and Vast.ai have Python APIs for programmatic access.
|
||||
|
||||
### Uninstalling
|
||||
|
||||
```bash
|
||||
rm -rf ~/.video-toolkit/propainter
|
||||
```
|
||||
|
||||
## Future Optional Components
|
||||
|
||||
The optional components system is designed to support additional ML-based tools:
|
||||
|
||||
- **Video upscaling** (Real-ESRGAN, etc.)
|
||||
- **Audio enhancement** (noise removal, etc.)
|
||||
- **Scene detection** (automatic scene splitting)
|
||||
|
||||
These will follow the same pattern: install on first use, stored in `~/.video-toolkit/`.
|
||||
@@ -0,0 +1,312 @@
|
||||
# RunPod Cloud GPU Setup
|
||||
|
||||
This guide covers setting up RunPod serverless GPUs for watermark removal (and future GPU-intensive video tools).
|
||||
|
||||
## Why RunPod?
|
||||
|
||||
The dewatermark tool uses ProPainter, an AI inpainting model that requires significant GPU power:
|
||||
|
||||
| Hardware | Processing Time (30s video) | Viable? |
|
||||
|----------|----------------------------|---------|
|
||||
| NVIDIA RTX 3090 | 2-5 minutes | Yes |
|
||||
| Apple Silicon M1/M2/M3 | 4+ hours | No |
|
||||
| CPU only | 10+ hours | No |
|
||||
|
||||
RunPod provides on-demand NVIDIA GPUs at ~$0.34/hour, making it cost-effective for occasional use (~$0.05-0.30 per video).
|
||||
|
||||
## Quick Start (Automated)
|
||||
|
||||
The fastest way to set up RunPod:
|
||||
|
||||
```bash
|
||||
# 1. Add your RunPod API key to .env
|
||||
echo "RUNPOD_API_KEY=your_key_here" >> .env
|
||||
|
||||
# 2. Run automated setup (creates template + endpoint)
|
||||
python tools/dewatermark.py --setup
|
||||
|
||||
# 3. Done! Now use it:
|
||||
python tools/dewatermark.py --input video.mp4 --region x,y,w,h --output out.mp4 --runpod
|
||||
```
|
||||
|
||||
The `--setup` command will:
|
||||
- Create a serverless template using the public Docker image
|
||||
- Create an endpoint with RTX 3090 GPU (AMPERE_24)
|
||||
- Save the endpoint ID to your `.env` file
|
||||
|
||||
Use `--setup-gpu AMPERE_16` for RTX 3080 or `--setup-gpu ADA_24` for RTX 4090.
|
||||
|
||||
---
|
||||
|
||||
## Manual Setup
|
||||
|
||||
If you prefer to set up manually via the web console:
|
||||
|
||||
### 1. Create RunPod Account
|
||||
|
||||
1. Go to [runpod.io](https://runpod.io) and sign up
|
||||
2. Add credits to your account ($10 minimum, lasts for many videos)
|
||||
3. Go to Settings > API Keys and create an API key
|
||||
|
||||
### 2. Create Serverless Endpoint
|
||||
|
||||
A pre-built public image is available:
|
||||
```
|
||||
ghcr.io/conalmullan/video-toolkit-propainter:v2.0.0
|
||||
```
|
||||
|
||||
> **Note:** Use versioned tags (not `:latest`) to ensure workers pull the correct image.
|
||||
|
||||
Alternatively, build your own (see `docker/runpod-propainter/README.md`).
|
||||
|
||||
1. Go to [RunPod Serverless Console](https://www.runpod.io/console/serverless)
|
||||
2. Click **New Endpoint**
|
||||
3. Configure:
|
||||
|
||||
| Setting | Value | Notes |
|
||||
|---------|-------|-------|
|
||||
| Docker Image | `ghcr.io/conalmullan/video-toolkit-propainter:latest` | Public image |
|
||||
| GPU | RTX 3090 or RTX 4090 | 24GB VRAM recommended |
|
||||
| Max Workers | 1 | Scale up if processing many videos |
|
||||
| Idle Timeout | 5 seconds | Fast scale-down to save costs |
|
||||
| Execution Timeout | 3600 seconds | 1 hour max per job |
|
||||
|
||||
4. Click **Create Endpoint**
|
||||
5. Copy the **Endpoint ID** (looks like: `abc123xyz`)
|
||||
|
||||
### 3. Configure Local Environment
|
||||
|
||||
Add to your `.env` file:
|
||||
|
||||
```bash
|
||||
# RunPod Configuration
|
||||
RUNPOD_API_KEY=your_api_key_here
|
||||
RUNPOD_ENDPOINT_ID=your_endpoint_id_here
|
||||
```
|
||||
|
||||
### 4. Test It
|
||||
|
||||
```bash
|
||||
# Dry run (doesn't actually process)
|
||||
python tools/dewatermark.py \
|
||||
--input video.mp4 \
|
||||
--region 1080,660,195,40 \
|
||||
--output clean.mp4 \
|
||||
--runpod \
|
||||
--dry-run
|
||||
|
||||
# Real processing
|
||||
python tools/dewatermark.py \
|
||||
--input video.mp4 \
|
||||
--region 1080,660,195,40 \
|
||||
--output clean.mp4 \
|
||||
--runpod
|
||||
```
|
||||
|
||||
## How It Works
|
||||
|
||||
```
|
||||
1. Local tool uploads video to temporary storage
|
||||
2. Submits job to RunPod endpoint
|
||||
3. RunPod spins up GPU worker (~30s cold start)
|
||||
4. Worker downloads video, runs ProPainter
|
||||
5. Worker uploads result, returns URL
|
||||
6. Local tool downloads result
|
||||
7. Worker scales down (you stop paying)
|
||||
```
|
||||
|
||||
## Cost Breakdown
|
||||
|
||||
### Per-Video Costs
|
||||
|
||||
| Video Length | Processing Time | Cost (RTX 3090) |
|
||||
|--------------|-----------------|-----------------|
|
||||
| < 30 seconds | 2-5 minutes | ~$0.02 |
|
||||
| 30s - 2 min | 5-15 minutes | ~$0.08 |
|
||||
| 2 - 5 min | 15-45 minutes | ~$0.25 |
|
||||
| > 5 min | 45+ minutes | ~$0.40+ |
|
||||
|
||||
### GPU Options
|
||||
|
||||
| GPU | VRAM | Cost/hr | Speed | Best For |
|
||||
|-----|------|---------|-------|----------|
|
||||
| RTX 3090 | 24GB | $0.34 | Fast | Most videos (recommended) |
|
||||
| RTX 4090 | 24GB | $0.69 | Faster | Tight deadlines |
|
||||
| A100 | 80GB | $1.99 | Fastest | Very long videos |
|
||||
|
||||
### Tips to Minimize Costs
|
||||
|
||||
1. **Use 5-second idle timeout** - Workers scale down quickly
|
||||
2. **Process in batches** - Submit multiple videos to same warm worker
|
||||
3. **Right-size your GPU** - RTX 3090 is plenty for most videos
|
||||
4. **Set max workers = 1** initially - Prevents runaway costs
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### "RUNPOD_API_KEY not set"
|
||||
|
||||
Add your API key to `.env`:
|
||||
```bash
|
||||
RUNPOD_API_KEY=your_key_here
|
||||
```
|
||||
|
||||
### "RUNPOD_ENDPOINT_ID not set"
|
||||
|
||||
Add your endpoint ID to `.env`:
|
||||
```bash
|
||||
RUNPOD_ENDPOINT_ID=abc123xyz
|
||||
```
|
||||
|
||||
### Job times out
|
||||
|
||||
Default timeout is 30 minutes. For longer videos:
|
||||
```bash
|
||||
python tools/dewatermark.py ... --runpod --runpod-timeout 3600
|
||||
```
|
||||
|
||||
### "Failed to upload video"
|
||||
|
||||
- Check your internet connection
|
||||
- Verify the video file exists and is readable
|
||||
- Large files (>500MB) may take several minutes to upload
|
||||
|
||||
### Cold start is slow (~30-60 seconds)
|
||||
|
||||
This is normal for the first request after idle. The worker needs to:
|
||||
1. Spin up the container
|
||||
2. Load PyTorch and models into GPU memory
|
||||
|
||||
Subsequent requests to a warm worker are faster.
|
||||
|
||||
### "ProPainter processing failed"
|
||||
|
||||
Check the RunPod logs:
|
||||
1. Go to RunPod Console > Serverless > Your Endpoint > Logs
|
||||
2. Look for error messages from the handler
|
||||
|
||||
Common issues:
|
||||
- Video format not supported (try converting to MP4)
|
||||
- Region coordinates exceed video dimensions
|
||||
- GPU ran out of memory (shouldn't happen with 24GB GPUs)
|
||||
|
||||
## File Transfer: Cloudflare R2 (Recommended)
|
||||
|
||||
By default, videos are uploaded via free file hosting services (litterbox.catbox.moe, etc.). These work but can be unreliable for large files.
|
||||
|
||||
**Cloudflare R2** provides reliable, fast file transfer with a generous free tier:
|
||||
- **10 GB storage** (we clean up after each job)
|
||||
- **10 million operations/month**
|
||||
- **Zero egress fees** (unlike AWS S3)
|
||||
- **No expiration** (unlike AWS's 12-month free tier)
|
||||
|
||||
### R2 Setup
|
||||
|
||||
1. **Create Cloudflare Account** (free): https://dash.cloudflare.com
|
||||
|
||||
2. **Create R2 Bucket**:
|
||||
- Go to R2 Object Storage → Create bucket
|
||||
- Name: `video-toolkit` (or any name)
|
||||
- Click Create
|
||||
|
||||
3. **Create API Token**:
|
||||
- R2 → Overview → Manage R2 API Tokens
|
||||
- Create API Token → Object Read & Write
|
||||
- Specify bucket: `video-toolkit`
|
||||
- Copy the **Access Key ID** and **Secret Access Key** (shown once!)
|
||||
|
||||
4. **Get Account ID**:
|
||||
- Visible in dashboard URL: `dash.cloudflare.com/<ACCOUNT_ID>/r2`
|
||||
|
||||
5. **Add to .env**:
|
||||
```bash
|
||||
R2_ACCOUNT_ID=your_account_id
|
||||
R2_ACCESS_KEY_ID=your_access_key_id
|
||||
R2_SECRET_ACCESS_KEY=your_secret_access_key
|
||||
R2_BUCKET_NAME=video-toolkit
|
||||
```
|
||||
|
||||
6. **Install boto3** (if not already):
|
||||
```bash
|
||||
pip install boto3
|
||||
```
|
||||
|
||||
That's it! The dewatermark tool will automatically use R2 for file transfer.
|
||||
|
||||
### Without R2
|
||||
|
||||
If R2 is not configured, the tool falls back to free file hosting services:
|
||||
- `litterbox.catbox.moe` (200MB, 24h retention)
|
||||
- `file.io` (2GB, 1 download)
|
||||
- `transfer.sh` (10GB, 14 days) - often down
|
||||
- `0x0.st` (512MB, 30 days) - blocks many requests
|
||||
|
||||
These work for testing but may fail intermittently for production use.
|
||||
|
||||
## Advanced Configuration
|
||||
|
||||
### Multiple Endpoints
|
||||
|
||||
You can create multiple endpoints for different use cases:
|
||||
|
||||
```bash
|
||||
# .env
|
||||
RUNPOD_ENDPOINT_ID=abc123xyz # Default (RTX 3090)
|
||||
RUNPOD_ENDPOINT_ID_FAST=def456uvw # Fast (RTX 4090)
|
||||
```
|
||||
|
||||
### Monitoring Usage
|
||||
|
||||
1. Go to RunPod Console > Usage
|
||||
2. View spend by endpoint, GPU type, and time period
|
||||
3. Set up billing alerts to avoid surprises
|
||||
|
||||
## Security Notes
|
||||
|
||||
- API keys grant full access to your RunPod account - keep them secret
|
||||
- R2 credentials are passed to RunPod workers for result upload - ensure your bucket is private
|
||||
- Without R2, videos go through public file hosting services (not recommended for sensitive content)
|
||||
- R2 objects are automatically cleaned up after download
|
||||
- Presigned URLs expire after 2 hours
|
||||
|
||||
## Future GPU Tools
|
||||
|
||||
The RunPod handler is designed for extensibility. Future operations may include:
|
||||
|
||||
- **upscale** - Video upscaling with Real-ESRGAN
|
||||
- **denoise** - Audio/video denoising
|
||||
- **stabilize** - Video stabilization
|
||||
- **style-transfer** - AI style transfer
|
||||
|
||||
These would use the same endpoint and Docker image, just with different `operation` values.
|
||||
|
||||
## Current Status & Known Limitations
|
||||
|
||||
**Working (as of v2.0.0):**
|
||||
- ✅ End-to-end watermark removal via RunPod
|
||||
- ✅ Cloudflare R2 file transfer (reliable, fast)
|
||||
- ✅ Automatic GPU detection (respects RunPod's CUDA_VISIBLE_DEVICES)
|
||||
- ✅ Smart auto resize_ratio based on VRAM + video size
|
||||
- ✅ 30-second video processing confirmed working
|
||||
|
||||
**Current Limitations:**
|
||||
|
||||
| Issue | Description | Workaround |
|
||||
|-------|-------------|------------|
|
||||
| Untested with long videos | Only 30-second clips tested | Try longer chunks in next session |
|
||||
| Full resolution OOM | `resize_ratio=1.0` may fail on GPUs with <48GB VRAM | Use `auto` (default) or upscale result post-processing |
|
||||
|
||||
**Next Steps (planned):**
|
||||
- [ ] Test longer video chunks (1-3 minutes)
|
||||
- [ ] Add post-processing upscale option to restore full resolution
|
||||
- [ ] Profile memory usage at different resolutions
|
||||
- [ ] Consider chunking very long videos client-side
|
||||
|
||||
## Version History
|
||||
|
||||
| Version | Date | Changes |
|
||||
|---------|------|---------|
|
||||
| v2.0.0 | 2025-12-30 | **Major fix:** GPU detection (removed CUDA_VISIBLE_DEVICES), CUDA 12.4, auto resize_ratio |
|
||||
| v1.2.1 | 2025-12-30 | Fix output file detection (`inpaint_out.mp4` not `masked_in.mp4`) |
|
||||
| v1.2.0 | 2025-12-30 | Fix GPU detection (use max across all GPUs), improve memory profiles |
|
||||
| v1.1.0 | 2025-12-30 | Add `resize_ratio` parameter, improve error logging |
|
||||
| v1.0.0 | 2025-12-30 | Initial R2 integration, NumPy 1.x fix |
|
||||
@@ -55,3 +55,45 @@ def get_default_output_dir(project_path: str | None = None) -> Path:
|
||||
if project_path:
|
||||
return Path(project_path) / "public" / "audio"
|
||||
return find_workspace_root() / "public" / "audio"
|
||||
|
||||
|
||||
def get_runpod_api_key() -> str | None:
|
||||
"""Get RunPod API key from environment."""
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
return os.getenv("RUNPOD_API_KEY")
|
||||
|
||||
|
||||
def get_runpod_endpoint_id() -> str | None:
|
||||
"""Get RunPod endpoint ID from environment."""
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
return os.getenv("RUNPOD_ENDPOINT_ID")
|
||||
|
||||
|
||||
def get_r2_config() -> dict | None:
|
||||
"""Get Cloudflare R2 configuration from environment.
|
||||
|
||||
Returns dict with account_id, access_key_id, secret_access_key, bucket_name
|
||||
or None if not configured.
|
||||
"""
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
|
||||
account_id = os.getenv("R2_ACCOUNT_ID")
|
||||
access_key_id = os.getenv("R2_ACCESS_KEY_ID")
|
||||
secret_access_key = os.getenv("R2_SECRET_ACCESS_KEY")
|
||||
bucket_name = os.getenv("R2_BUCKET_NAME", "video-toolkit")
|
||||
|
||||
# Check if R2 is configured (all required fields present and not placeholder)
|
||||
if (account_id and access_key_id and secret_access_key
|
||||
and account_id != "your_account_id_here"
|
||||
and access_key_id != "your_access_key_id_here"):
|
||||
return {
|
||||
"account_id": account_id,
|
||||
"access_key_id": access_key_id,
|
||||
"secret_access_key": secret_access_key,
|
||||
"bucket_name": bucket_name,
|
||||
"endpoint_url": f"https://{account_id}.r2.cloudflarestorage.com",
|
||||
}
|
||||
return None
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,560 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Locate and verify watermark positions in video files.
|
||||
|
||||
This tool helps identify watermark coordinates for use with dewatermark.py.
|
||||
It extracts frames, overlays grids, and marks regions for visual verification.
|
||||
|
||||
Usage:
|
||||
# Interactive exploration - extract frames with grid overlay
|
||||
python tools/locate_watermark.py --input video.mp4 --grid --output-dir /tmp/review/
|
||||
|
||||
# Verify a specific region across multiple frames
|
||||
python tools/locate_watermark.py --input video.mp4 --region 1100,650,150,50 --verify
|
||||
|
||||
# Use a preset for common watermarks
|
||||
python tools/locate_watermark.py --input video.mp4 --preset notebooklm --verify
|
||||
|
||||
# Quick check - mark single frame
|
||||
python tools/locate_watermark.py --input video.mp4 --region 1100,650,150,50 --mark
|
||||
|
||||
Presets:
|
||||
notebooklm - Bottom-right corner (Google NotebookLM videos)
|
||||
tiktok - Bottom-center username area
|
||||
stock-br - Bottom-right stock footage watermark
|
||||
stock-bl - Bottom-left stock footage watermark
|
||||
stock-center - Center watermark (common in stock footage)
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
# Watermark presets (x, y, width, height) - will be scaled to video dimensions
|
||||
PRESETS = {
|
||||
"notebooklm": {
|
||||
"description": "Google NotebookLM - bottom-right corner",
|
||||
"region_1280x720": (1100, 650, 150, 50),
|
||||
"region_1920x1080": (1650, 975, 225, 75),
|
||||
},
|
||||
"tiktok": {
|
||||
"description": "TikTok username - bottom-center",
|
||||
"region_1080x1920": (340, 1750, 400, 80), # Portrait
|
||||
"region_1280x720": (440, 650, 400, 50), # Landscape
|
||||
},
|
||||
"stock-br": {
|
||||
"description": "Stock footage - bottom-right",
|
||||
"region_1280x720": (1000, 620, 260, 80),
|
||||
"region_1920x1080": (1500, 930, 390, 120),
|
||||
},
|
||||
"stock-bl": {
|
||||
"description": "Stock footage - bottom-left",
|
||||
"region_1280x720": (20, 620, 260, 80),
|
||||
"region_1920x1080": (30, 930, 390, 120),
|
||||
},
|
||||
"stock-center": {
|
||||
"description": "Stock footage - center watermark",
|
||||
"region_1280x720": (440, 260, 400, 200),
|
||||
"region_1920x1080": (660, 390, 600, 300),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Locate and verify watermark positions in video",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Extract frames with coordinate grid for exploration
|
||||
python tools/locate_watermark.py --input video.mp4 --grid --output-dir ./review/
|
||||
|
||||
# Verify NotebookLM watermark position
|
||||
python tools/locate_watermark.py --input video.mp4 --preset notebooklm --verify
|
||||
|
||||
# Mark custom region on multiple frames
|
||||
python tools/locate_watermark.py --input video.mp4 --region 1100,650,150,50 --verify
|
||||
|
||||
# Output coordinates as JSON (for scripting)
|
||||
python tools/locate_watermark.py --input video.mp4 --preset notebooklm --json
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--input", "-i",
|
||||
type=str,
|
||||
help="Input video file path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--region", "-r",
|
||||
type=str,
|
||||
help="Watermark region as x,y,width,height (e.g., 1100,650,150,50)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preset", "-p",
|
||||
type=str,
|
||||
choices=list(PRESETS.keys()),
|
||||
help="Use a preset watermark position",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir", "-o",
|
||||
type=str,
|
||||
help="Directory to save marked frames (default: temp directory)",
|
||||
)
|
||||
|
||||
# Actions
|
||||
parser.add_argument(
|
||||
"--grid",
|
||||
action="store_true",
|
||||
help="Overlay coordinate grid on frames",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mark",
|
||||
action="store_true",
|
||||
help="Mark region with rectangle on frames",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verify",
|
||||
action="store_true",
|
||||
help="Extract multiple frames and mark region for verification",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--crop",
|
||||
action="store_true",
|
||||
help="Also output cropped watermark regions",
|
||||
)
|
||||
|
||||
# Sampling options
|
||||
parser.add_argument(
|
||||
"--samples",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Number of frames to extract (default: 5)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--timestamps",
|
||||
type=str,
|
||||
help="Specific timestamps to extract (comma-separated, e.g., '10,30,60,90')",
|
||||
)
|
||||
|
||||
# Grid options
|
||||
parser.add_argument(
|
||||
"--grid-spacing",
|
||||
type=int,
|
||||
default=50,
|
||||
help="Grid line spacing in pixels (default: 50)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--grid-region",
|
||||
type=str,
|
||||
help="Only show grid in region x,y,width,height (default: bottom-right quadrant)",
|
||||
)
|
||||
|
||||
# Output options
|
||||
parser.add_argument(
|
||||
"--json",
|
||||
action="store_true",
|
||||
help="Output result as JSON",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--list-presets",
|
||||
action="store_true",
|
||||
help="List available watermark presets",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--open",
|
||||
action="store_true",
|
||||
help="Open output directory in Finder after processing (macOS)",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def get_video_info(video_path: str) -> dict | None:
|
||||
"""Get video dimensions and duration using ffprobe."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffprobe", "-v", "error",
|
||||
"-select_streams", "v:0",
|
||||
"-show_entries", "stream=width,height",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "json",
|
||||
video_path,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
data = json.loads(result.stdout)
|
||||
stream = data.get("streams", [{}])[0]
|
||||
fmt = data.get("format", {})
|
||||
return {
|
||||
"width": stream.get("width"),
|
||||
"height": stream.get("height"),
|
||||
"duration": float(fmt.get("duration", 0)),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def parse_region(region_str: str) -> tuple[int, int, int, int] | None:
|
||||
"""Parse region string 'x,y,width,height' into tuple."""
|
||||
try:
|
||||
parts = [int(x.strip()) for x in region_str.split(",")]
|
||||
if len(parts) == 4:
|
||||
return tuple(parts)
|
||||
except ValueError:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def get_preset_region(preset_name: str, width: int, height: int) -> tuple[int, int, int, int] | None:
|
||||
"""Get region for a preset, scaled to video dimensions."""
|
||||
if preset_name not in PRESETS:
|
||||
return None
|
||||
|
||||
preset = PRESETS[preset_name]
|
||||
key = f"region_{width}x{height}"
|
||||
|
||||
# Try exact match first
|
||||
if key in preset:
|
||||
return preset[key]
|
||||
|
||||
# Find closest match and scale
|
||||
for preset_key, region in preset.items():
|
||||
if preset_key.startswith("region_"):
|
||||
dims = preset_key.replace("region_", "").split("x")
|
||||
preset_w, preset_h = int(dims[0]), int(dims[1])
|
||||
|
||||
# Scale proportionally
|
||||
scale_x = width / preset_w
|
||||
scale_y = height / preset_h
|
||||
|
||||
x, y, w, h = region
|
||||
return (
|
||||
int(x * scale_x),
|
||||
int(y * scale_y),
|
||||
int(w * scale_x),
|
||||
int(h * scale_y),
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extract_frame(video_path: str, timestamp: float, output_path: str) -> bool:
|
||||
"""Extract a single frame from video."""
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffmpeg", "-y",
|
||||
"-ss", str(timestamp),
|
||||
"-i", video_path,
|
||||
"-frames:v", "1",
|
||||
output_path,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
return result.returncode == 0
|
||||
|
||||
|
||||
def add_grid_overlay(
|
||||
input_path: str,
|
||||
output_path: str,
|
||||
width: int,
|
||||
height: int,
|
||||
spacing: int = 50,
|
||||
region: tuple[int, int, int, int] | None = None,
|
||||
) -> bool:
|
||||
"""Add coordinate grid overlay to image using ImageMagick."""
|
||||
|
||||
# Determine grid region (default to bottom-right quadrant)
|
||||
if region:
|
||||
grid_x, grid_y, grid_w, grid_h = region
|
||||
else:
|
||||
# Bottom-right quadrant
|
||||
grid_x = width // 2
|
||||
grid_y = height // 2
|
||||
grid_w = width // 2
|
||||
grid_h = height // 2
|
||||
|
||||
draw_commands = []
|
||||
|
||||
# Vertical lines
|
||||
for x in range(grid_x, grid_x + grid_w + 1, spacing):
|
||||
if x <= width:
|
||||
draw_commands.append(f"line {x},{grid_y} {x},{min(grid_y + grid_h, height)}")
|
||||
# Label at bottom
|
||||
label_y = min(grid_y + grid_h - 5, height - 5)
|
||||
draw_commands.append(f"text {x+2},{label_y} '{x}'")
|
||||
|
||||
# Horizontal lines
|
||||
for y in range(grid_y, grid_y + grid_h + 1, spacing):
|
||||
if y <= height:
|
||||
draw_commands.append(f"line {grid_x},{y} {min(grid_x + grid_w, width)},{y}")
|
||||
# Label at left
|
||||
draw_commands.append(f"text {grid_x+2},{y-2} '{y}'")
|
||||
|
||||
cmd = [
|
||||
"magick", input_path,
|
||||
"-stroke", "yellow",
|
||||
"-strokewidth", "1",
|
||||
"-fill", "yellow",
|
||||
"-pointsize", "12",
|
||||
]
|
||||
|
||||
for draw_cmd in draw_commands:
|
||||
cmd.extend(["-draw", draw_cmd])
|
||||
|
||||
cmd.append(output_path)
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
return result.returncode == 0
|
||||
|
||||
|
||||
def mark_region(
|
||||
input_path: str,
|
||||
output_path: str,
|
||||
region: tuple[int, int, int, int],
|
||||
color: str = "red",
|
||||
stroke_width: int = 3,
|
||||
) -> bool:
|
||||
"""Mark a region with a rectangle using ImageMagick."""
|
||||
x, y, w, h = region
|
||||
x2, y2 = x + w, y + h
|
||||
|
||||
result = subprocess.run(
|
||||
[
|
||||
"magick", input_path,
|
||||
"-stroke", color,
|
||||
"-strokewidth", str(stroke_width),
|
||||
"-fill", "none",
|
||||
"-draw", f"rectangle {x},{y} {x2},{y2}",
|
||||
output_path,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
return result.returncode == 0
|
||||
|
||||
|
||||
def crop_region(
|
||||
input_path: str,
|
||||
output_path: str,
|
||||
region: tuple[int, int, int, int],
|
||||
) -> bool:
|
||||
"""Crop image to specified region."""
|
||||
x, y, w, h = region
|
||||
|
||||
result = subprocess.run(
|
||||
[
|
||||
"magick", input_path,
|
||||
"-crop", f"{w}x{h}+{x}+{y}",
|
||||
"+repage",
|
||||
output_path,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
return result.returncode == 0
|
||||
|
||||
|
||||
def calculate_timestamps(duration: float, num_samples: int, margin: float = 5.0) -> list[float]:
|
||||
"""Calculate evenly spaced timestamps across video duration."""
|
||||
# Avoid very start and end of video
|
||||
start = min(margin, duration * 0.05)
|
||||
end = max(duration - margin, duration * 0.95)
|
||||
|
||||
if num_samples == 1:
|
||||
return [duration / 2]
|
||||
|
||||
step = (end - start) / (num_samples - 1)
|
||||
return [start + i * step for i in range(num_samples)]
|
||||
|
||||
|
||||
def list_presets():
|
||||
"""Print available presets."""
|
||||
print("Available watermark presets:")
|
||||
print("-" * 50)
|
||||
for name, preset in PRESETS.items():
|
||||
print(f"\n {name}")
|
||||
print(f" {preset['description']}")
|
||||
for key, value in preset.items():
|
||||
if key.startswith("region_"):
|
||||
dims = key.replace("region_", "")
|
||||
print(f" {dims}: x={value[0]}, y={value[1]}, w={value[2]}, h={value[3]}")
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
# Handle --list-presets
|
||||
if args.list_presets:
|
||||
list_presets()
|
||||
return
|
||||
|
||||
# Check input is provided for other operations
|
||||
if not args.input:
|
||||
print("Error: --input is required", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Check input file
|
||||
if not Path(args.input).exists():
|
||||
print(f"Error: Input file not found: {args.input}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Check for ImageMagick
|
||||
if shutil.which("magick") is None:
|
||||
print("Error: ImageMagick not found. Install with: brew install imagemagick", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Get video info
|
||||
video_info = get_video_info(args.input)
|
||||
if not video_info:
|
||||
print("Error: Could not read video info", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
width = video_info["width"]
|
||||
height = video_info["height"]
|
||||
duration = video_info["duration"]
|
||||
|
||||
verbose = not args.json
|
||||
|
||||
if verbose:
|
||||
print(f"Video: {args.input}")
|
||||
print(f"Dimensions: {width}x{height}")
|
||||
print(f"Duration: {duration:.1f}s")
|
||||
|
||||
# Determine region
|
||||
region = None
|
||||
if args.region:
|
||||
region = parse_region(args.region)
|
||||
if not region:
|
||||
print(f"Error: Invalid region format: {args.region}", file=sys.stderr)
|
||||
print("Expected format: x,y,width,height (e.g., 1100,650,150,50)", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
elif args.preset:
|
||||
region = get_preset_region(args.preset, width, height)
|
||||
if verbose:
|
||||
print(f"Preset '{args.preset}': {region[0]},{region[1]},{region[2]},{region[3]}")
|
||||
|
||||
# Determine timestamps
|
||||
if args.timestamps:
|
||||
timestamps = [float(t.strip()) for t in args.timestamps.split(",")]
|
||||
else:
|
||||
timestamps = calculate_timestamps(duration, args.samples)
|
||||
|
||||
# Set up output directory
|
||||
if args.output_dir:
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
cleanup_temp = False
|
||||
else:
|
||||
output_dir = Path(tempfile.mkdtemp(prefix="locate_wm_"))
|
||||
cleanup_temp = not args.open # Keep if opening in Finder
|
||||
|
||||
if verbose:
|
||||
print(f"Output directory: {output_dir}")
|
||||
print()
|
||||
|
||||
# Process frames
|
||||
results = []
|
||||
for i, ts in enumerate(timestamps):
|
||||
if verbose:
|
||||
print(f"Processing frame {i+1}/{len(timestamps)} at {ts:.1f}s...")
|
||||
|
||||
# Extract frame
|
||||
frame_path = output_dir / f"frame_{ts:.0f}s.png"
|
||||
if not extract_frame(args.input, ts, str(frame_path)):
|
||||
print(f" Warning: Failed to extract frame at {ts}s", file=sys.stderr)
|
||||
continue
|
||||
|
||||
frame_result = {
|
||||
"timestamp": ts,
|
||||
"frame": str(frame_path),
|
||||
}
|
||||
|
||||
# Add grid overlay
|
||||
if args.grid:
|
||||
grid_region = None
|
||||
if args.grid_region:
|
||||
grid_region = parse_region(args.grid_region)
|
||||
|
||||
grid_path = output_dir / f"frame_{ts:.0f}s_grid.png"
|
||||
if add_grid_overlay(str(frame_path), str(grid_path), width, height, args.grid_spacing, grid_region):
|
||||
frame_result["grid"] = str(grid_path)
|
||||
if verbose:
|
||||
print(f" Created: {grid_path.name}")
|
||||
|
||||
# Mark region
|
||||
if (args.mark or args.verify) and region:
|
||||
marked_path = output_dir / f"frame_{ts:.0f}s_marked.png"
|
||||
source = frame_result.get("grid", str(frame_path))
|
||||
if mark_region(source, str(marked_path), region):
|
||||
frame_result["marked"] = str(marked_path)
|
||||
if verbose:
|
||||
print(f" Created: {marked_path.name}")
|
||||
|
||||
# Crop region
|
||||
if args.crop and region:
|
||||
crop_path = output_dir / f"frame_{ts:.0f}s_crop.png"
|
||||
if crop_region(str(frame_path), str(crop_path), region):
|
||||
frame_result["crop"] = str(crop_path)
|
||||
if verbose:
|
||||
print(f" Created: {crop_path.name}")
|
||||
|
||||
results.append(frame_result)
|
||||
|
||||
# Output
|
||||
output = {
|
||||
"input": args.input,
|
||||
"dimensions": f"{width}x{height}",
|
||||
"duration": duration,
|
||||
"region": f"{region[0]},{region[1]},{region[2]},{region[3]}" if region else None,
|
||||
"preset": args.preset,
|
||||
"output_dir": str(output_dir),
|
||||
"frames": results,
|
||||
}
|
||||
|
||||
if region:
|
||||
output["dewatermark_command"] = (
|
||||
f"python tools/dewatermark.py --input \"{args.input}\" "
|
||||
f"--region {region[0]},{region[1]},{region[2]},{region[3]} "
|
||||
f"--output \"output_clean.mp4\""
|
||||
)
|
||||
|
||||
if args.json:
|
||||
print(json.dumps(output, indent=2))
|
||||
else:
|
||||
print()
|
||||
print("=" * 50)
|
||||
print(f"Extracted {len(results)} frames to: {output_dir}")
|
||||
if region:
|
||||
print(f"Region: {region[0]},{region[1]},{region[2]},{region[3]}")
|
||||
print()
|
||||
print("To remove watermark, run:")
|
||||
print(f" python tools/dewatermark.py \\")
|
||||
print(f" --input \"{args.input}\" \\")
|
||||
print(f" --region {region[0]},{region[1]},{region[2]},{region[3]} \\")
|
||||
print(f" --output \"output_clean.mp4\"")
|
||||
print("=" * 50)
|
||||
|
||||
# Open in Finder (macOS)
|
||||
if args.open:
|
||||
subprocess.run(["open", str(output_dir)])
|
||||
|
||||
# Cleanup temp directory if not needed
|
||||
if cleanup_temp and not args.output_dir:
|
||||
# Don't cleanup - let user review
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,3 +1,5 @@
|
||||
# Video Toolkit Python Dependencies
|
||||
elevenlabs>=1.0.0
|
||||
python-dotenv>=1.0.0
|
||||
requests>=2.28.0
|
||||
boto3>=1.28.0 # For Cloudflare R2 (S3-compatible)
|
||||
|
||||
Reference in New Issue
Block a user