Showing preview only (614K chars total). Download the full file or copy to clipboard to get everything.
Repository: nikopueringer/CorridorKey
Branch: main
Commit: 3438680c4ae5
Files: 92
Total size: 584.6 KB
Directory structure:
gitextract_m8ktm8q9/
├── .dockerignore
├── .git-blame-ignore-revs
├── .github/
│ ├── ISSUE_TEMPLATE/
│ │ └── bug_report.yml
│ ├── pull_request_template.md
│ └── workflows/
│ ├── ci.yml
│ └── docs.yml
├── .gitignore
├── .python-version
├── BiRefNetModule/
│ ├── checkpoints/
│ │ └── .gitkeep
│ └── wrapper.py
├── CONTRIBUTING.md
├── ClipsForInference/
│ └── .gitkeep
├── CorridorKeyModule/
│ ├── IgnoredCheckpoints/
│ │ └── .gitkeep
│ ├── README.md
│ ├── __init__.py
│ ├── backend.py
│ ├── checkpoints/
│ │ └── .gitkeep
│ ├── core/
│ │ ├── __init__.py
│ │ ├── color_utils.py
│ │ └── model_transformer.py
│ └── inference_engine.py
├── CorridorKey_DRAG_CLIPS_HERE_local.bat
├── CorridorKey_DRAG_CLIPS_HERE_local.sh
├── Dockerfile
├── IgnoredClips/
│ └── .gitkeep
├── Install_CorridorKey_Linux_Mac.sh
├── Install_CorridorKey_Windows.bat
├── Install_GVM_Linux_Mac.sh
├── Install_GVM_Windows.bat
├── Install_VideoMaMa_Linux_Mac.sh
├── Install_VideoMaMa_Windows.bat
├── LICENSE
├── Output/
│ └── .gitkeep
├── README.md
├── RunGVMOnly.sh
├── RunInferenceOnly.sh
├── VideoMaMaInferenceModule/
│ ├── LICENSE.md
│ ├── README.md
│ ├── __init__.py
│ ├── checkpoints/
│ │ └── .gitkeep
│ ├── inference.py
│ └── pipeline.py
├── backend/
│ ├── __init__.py
│ ├── clip_state.py
│ ├── errors.py
│ ├── ffmpeg_tools.py
│ ├── frame_io.py
│ ├── job_queue.py
│ ├── natural_sort.py
│ ├── project.py
│ ├── service.py
│ └── validators.py
├── clip_manager.py
├── corridorkey_cli.py
├── device_utils.py
├── docker-compose.yml
├── docs/
│ ├── LLM_HANDOVER.md
│ └── index.md
├── gvm_core/
│ ├── LICENSE.md
│ ├── README.md
│ ├── __init__.py
│ ├── gvm/
│ │ ├── __init__.py
│ │ ├── models/
│ │ │ ├── __init__.py
│ │ │ └── unet_spatio_temporal_condition.py
│ │ ├── pipelines/
│ │ │ └── pipeline_gvm.py
│ │ └── utils/
│ │ ├── __init__.py
│ │ └── inference_utils.py
│ ├── weights/
│ │ └── .gitkeep
│ └── wrapper.py
├── pyproject.toml
├── renovate.json
├── test_vram.py
├── tests/
│ ├── __init__.py
│ ├── conftest.py
│ ├── test_backend.py
│ ├── test_cli.py
│ ├── test_clip_manager.py
│ ├── test_color_utils.py
│ ├── test_device_utils.py
│ ├── test_e2e_workflow.py
│ ├── test_exr_gamma_bug_condition.py
│ ├── test_exr_gamma_preservation.py
│ ├── test_frame_io.py
│ ├── test_gamma_consistency.py
│ ├── test_imports.py
│ ├── test_inference_engine.py
│ ├── test_mlx_smoke.py
│ ├── test_pbt_auto_download.py
│ ├── test_pbt_backend_resolution.py
│ ├── test_pbt_dep_preservation.py
│ └── test_pyproject_structure.py
└── zensical.toml
================================================
FILE CONTENTS
================================================
================================================
FILE: .dockerignore
================================================
.git
.github
.pytest_cache
.ruff_cache
.uv-cache
.venv
__pycache__/
*.pyc
*.pyo
*.pyd
# Docker files (no need to copy themselves into the image)
Dockerfile
docker-compose.yml
# Documentation and repo metadata
CONTRIBUTING.md
docs/
.git-blame-ignore-revs
# Tests (not needed at runtime)
tests/
test_vram.py
# Launcher scripts (host-only)
*.bat
*.sh
# Large/generated data and local outputs
ClipsForInference/
IgnoredClips/
Output/
CorridorKeyModule/checkpoints/
gvm_core/weights/
VideoMaMaInferenceModule/checkpoints/
================================================
FILE: .git-blame-ignore-revs
================================================
# Automated code formatting — no behavioral changes
# ruff format + lint fixes
b0ad00efbc791ed097cd3fd241c10319beb8a631
================================================
FILE: .github/ISSUE_TEMPLATE/bug_report.yml
================================================
name: 🐛 Bug Report
description: Report a reproducible issue or unexpected behavior in the project.
title: "[Bug]: "
labels: [bug]
body:
- type: markdown
attributes:
value: |
Thank you for reporting a bug! Detailed bug reports help us fix issues faster.
**Before submitting:**
- Search existing issues to avoid duplicates
- Verify the issue is reproducible with the latest version
- Isolate the problem (provide minimal steps to reproduce)
- type: input
id: os
attributes:
label: Operating System
description: Your operating system and version.
placeholder: "Windows 11, macOS 14.2, Ubuntu 22.04 LTS"
validations:
required: true
- type: input
id: installation_method
attributes:
label: Installation Method
description: How did you install CorridorKey? (e.g., Windows batch installer, `uv sync`, Docker, manual setup)
placeholder: "Windows batch installer, uv sync --extra cuda, Docker"
validations:
required: true
- type: input
id: gpu_info
attributes:
label: GPU/Hardware (Optional)
description: GPU model and VRAM if applicable, or other relevant hardware constraints.
placeholder: "NVIDIA RTX 4090 (24GB), Apple M3 Pro (18GB unified memory)"
validations:
required: true
- type: textarea
id: steps
attributes:
label: Steps to Reproduce
description: Provide clear, numbered steps to reproduce the issue. Be as specific as possible.
placeholder: |
1. Step 1
2. Step 2
3. Step 3
4. Step 4
5. Error occurs
validations:
required: true
- type: textarea
id: expected
attributes:
label: Expected Behavior
description: What should happen instead? Be clear and concise.
placeholder: "The application should display a settings panel without errors."
validations:
required: false
- type: textarea
id: actual
attributes:
label: Actual Behavior
description: What happened instead? Describe the bug clearly and concisely.
placeholder: "The application crashes with an Error"
validations:
required: true
- type: textarea
id: logs
attributes:
label: Relevant Logs or Error Messages
description: Paste complete error logs, full stack traces, or screenshots. Use code blocks for readability.
placeholder: |
```
Put your log in here
```
- type: textarea
id: workaround
attributes:
label: Workaround (if available)
description: If you've found a way to work around this issue, please describe it.
placeholder: "As a workaround, I can set the timeout environment variable before starting the app."
validations:
required: false
- type: checkboxes
id: checks
attributes:
label: Verification Checklist
options:
- label: I've verified this bug hasn't been reported before
required: true
- label: I can reproduce this issue consistently
required: true
- label: I've included all relevant logs, screenshots, and error messages
required: true
- label: I've tested with the latest version of the project
required: false
================================================
FILE: .github/pull_request_template.md
================================================
## What does this change?
## How was it tested?
## Checklist
- [ ] `uv run pytest` passes
- [ ] `uv run ruff check` passes
- [ ] `uv run ruff format --check` passes
================================================
FILE: .github/workflows/ci.yml
================================================
name: CI
on:
push:
branches: [main]
pull_request:
branches: [main]
env:
UV_NO_SYNC: 1
UV_LOCKED: 1
OPENCV_IO_ENABLE_OPENEXR: 1
permissions:
contents: read
checks: write
pull-requests: write
jobs:
lint:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
with:
persist-credentials: false
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
- name: Install dependencies
run: uv sync --group dev
- name: Check formatting
run: uv run ruff format --check
- name: Check lint
run: uv run ruff check
test:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.10", "3.13"]
steps:
- uses: actions/checkout@v6
with:
persist-credentials: false
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
- name: Install dependencies
run: uv sync --group dev --python ${{ matrix.python-version }}
- name: Run tests
run: uv run pytest -v --tb=short -m "not gpu"
================================================
FILE: .github/workflows/docs.yml
================================================
name: Documentation
on:
push:
branches:
- main
paths:
- "docs/**"
- "zensical.toml"
workflow_dispatch:
permissions:
contents: read
pages: write
id-token: write
jobs:
deploy:
environment:
name: github-pages
url: ${{ steps.deployment.outputs.page_url }}
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: actions/configure-pages@v5
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
- name: Install docs dependencies via uv
run: uv sync --locked --only-group docs
- name: Build docs with zensical via uv
run: uv run zensical build --clean
- uses: actions/upload-pages-artifact@v4
with:
path: site
- uses: actions/deploy-pages@v4
id: deployment
================================================
FILE: .gitignore
================================================
# Python
__pycache__/
*.pyc
*.pyo
*.pyd
.coverage
.Python
.pytest_cache/
.venv/
.hypothesis
env/
venv/
ENV/
env.bak/
venv.bak/
# Project Specific
ClipsForInference/*
!ClipsForInference/.gitkeep
Output/*
!Output/.gitkeep
Ignored*/*
!Ignored*/.gitkeep
CorridorKey_remote.bat
.ipynb_checkpoints/
.DS_Store
# IDE
.vscode/
.idea/
# Models & Checkpoints (Large Files)
*.pth
*.pt
*.ckpt
*.safetensors
*.bin
*.onnx
# Checkpoint Directories
CorridorKeyModule/checkpoints/*
!CorridorKeyModule/checkpoints/.gitkeep
CorridorKeyModule/IgnoredCheckpoints/*
!CorridorKeyModule/IgnoredCheckpoints/.gitkeep
VideoMaMaInferenceModule/checkpoints/*
!VideoMaMaInferenceModule/checkpoints/.gitkeep
gvm_core/weights/*
!gvm_core/weights/.gitkeep
BiRefNetModule/checkpoints/*
!BiRefNetModule/checkpoints/.gitkeep
site
================================================
FILE: .python-version
================================================
3.13
================================================
FILE: BiRefNetModule/checkpoints/.gitkeep
================================================
================================================
FILE: BiRefNetModule/wrapper.py
================================================
import logging
import os
from pathlib import Path
from typing import Tuple
import cv2
import numpy as np
import torch
from huggingface_hub import snapshot_download
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
torch.set_float32_matmul_precision(["high", "highest"][0])
class ImagePreprocessor:
def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
self.transform_image = transforms.Compose(
[
transforms.Resize(resolution),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
def proc(self, image: Image.Image) -> torch.Tensor:
image = self.transform_image(image)
return image
usage_to_weights_file = {
"General": "BiRefNet",
"General-dynamic": "BiRefNet_dynamic",
"General-HR": "BiRefNet_HR",
"General-Lite": "BiRefNet_lite",
"General-Lite-2K": "BiRefNet_lite-2K",
"General-reso_512": "BiRefNet_512x512",
"Matting": "BiRefNet-matting",
"Matting-dynamic": "BiRefNet_dynamic-matting",
"Matting-HR": "BiRefNet_HR-Matting",
"Matting-Lite": "BiRefNet_lite-matting",
"Portrait": "BiRefNet-portrait",
"DIS": "BiRefNet-DIS5K",
"HRSOD": "BiRefNet-HRSOD",
"COD": "BiRefNet-COD",
"DIS-TR_TEs": "BiRefNet-DIS5K-TR_TEs",
"General-legacy": "BiRefNet-legacy",
}
half_precision = True
base_folder = os.path.join(os.path.dirname(__file__), "checkpoints")
class BiRefNetHandler:
def __init__(self, device="cpu", usage="General"):
self.device = device
# Set resolution
if usage in ["General-Lite-2K"]:
self.resolution = (2560, 1440)
elif usage in ["General-reso_512"]:
self.resolution = (512, 512)
elif usage in ["General-HR", "Matting-HR"]:
self.resolution = (2048, 2048)
else:
if "-dynamic" in usage:
self.resolution = None
else:
self.resolution = (1024, 1024)
repo_name = usage_to_weights_file[usage]
repo_id = f"ZhengPeng7/{repo_name}"
model_local_dir = os.path.join(base_folder, repo_name)
snapshot_download(
repo_id=repo_id,
local_dir=model_local_dir,
local_dir_use_symlinks=False, # Ensures actual files are downloaded, not just symlinks to the cache
)
self.birefnet = AutoModelForImageSegmentation.from_pretrained(model_local_dir, trust_remote_code=True)
self.birefnet.to(device)
self.birefnet.eval()
if half_precision:
self.birefnet.half()
def cleanup(self):
"""Explicitly clear model and release GPU memory."""
# Delete the model reference
if hasattr(self, "birefnet"):
del self.birefnet
# Clear Python garbage
import gc
gc.collect()
# Clear PyTorch CUDA cache
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
def process(self, input_path, alpha_output_dir=None, dilate_radius=0, on_frame_complete=None):
"""
Process a single video or directory of images.
"""
input_path = Path(input_path)
file_name = input_path.stem
is_video = input_path.suffix.lower() in [".mp4", ".mkv", ".gif", ".mov", ".avi"]
def get_frames():
"""Yields tuples of (image_numpy_array, output_file_name)"""
if is_video:
cap = cv2.VideoCapture(str(input_path))
count = 0
while True:
success, img = cap.read()
if not success:
break
yield img, f"{file_name}_alpha_{count:05d}.png"
count += 1
cap.release()
else:
image_files = sorted(
[
f
for f in input_path.iterdir()
if f.is_file() and f.suffix.lower() in [".jpg", ".png", ".jpeg", ".exr"]
]
)
if not image_files:
logging.warning(f"No images found in {input_path}")
return
# Setup EXR support once if needed
if "OPENCV_IO_ENABLE_OPENEXR" not in os.environ:
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
for img_path in image_files:
img = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)
if img is None:
continue
# Keep original filename for image sequences
yield img, f"alphaSeq_{img_path.stem}.png"
count = 0
for image, out_name in get_frames():
# Ensure correct conversion to RGB regardless of input format (EXR/PNG/JPG)
if len(image.shape) == 2:
image_rgb = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
elif image.shape[2] == 4:
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
else:
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# EXR images load as float32. PIL expects uint8. Normalize if necessary.
if image_rgb.dtype != np.uint8:
image_rgb = cv2.normalize(image_rgb, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
pil_image = Image.fromarray(image_rgb)
# Preprocess
if self.resolution is None: # Account for dynamic models
resolution_div_by_32 = [int(int(reso) // 32 * 32) for reso in pil_image.size]
if resolution_div_by_32 != self.resolution:
self.resolution = resolution_div_by_32
image_preprocessor = ImagePreprocessor(resolution=tuple(self.resolution))
image_proc = image_preprocessor.proc(pil_image).unsqueeze(0).to(self.device)
if half_precision:
image_proc = image_proc.half()
# Inference
with torch.no_grad():
preds = self.birefnet(image_proc)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred.float())
# Post-Process
target_size = (image.shape[1], image.shape[0])
mask = pred_pil.resize(target_size)
mask_np = np.array(mask)
# Dilate
if dilate_radius != 0:
abs_radius = abs(dilate_radius)
k_size = abs_radius * 2 + 1
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k_size, k_size))
if dilate_radius > 0:
mask_np = cv2.dilate(mask_np, kernel, iterations=1) # Expansion
else:
mask_np = cv2.erode(mask_np, kernel, iterations=1) # Contraction
# Strict Binary Threshold
_, mask_np = cv2.threshold(mask_np, 10, 255, cv2.THRESH_BINARY)
# Save
if alpha_output_dir:
save_path = os.path.join(alpha_output_dir, out_name)
cv2.imwrite(save_path, mask_np)
if on_frame_complete:
on_frame_complete(count, 0)
================================================
FILE: CONTRIBUTING.md
================================================
# Contributing to CorridorKey
Thanks for your interest in improving CorridorKey! Whether you're a VFX artist, a pipeline TD, or a machine learning researcher, contributions of all kinds are welcome — bug reports, feature ideas, documentation fixes, and code.
## Legal Agreement
By contributing to this project, you agree that your contributions will be licensed under the project's **[CorridorKey Licence](LICENSE)**.
By submitting a Pull Request, you specifically acknowledge and agree to the terms set forth in **Section 6 (CONTRIBUTIONS)** of the license. This ensures that Corridor Digital maintains the full right to use, distribute, and sublicense this codebase, including PR contributions. This is a project for the community, and will always remain freely available here.
## Getting Started
### Prerequisites
- Python 3.10 or newer
- [uv](https://docs.astral.sh/uv/) for dependency management
### Dev Setup
```bash
git clone https://github.com/nikopueringer/CorridorKey.git
cd CorridorKey
uv sync --group dev # installs all dependencies + dev tools (pytest, ruff)
```
That's it. No manual virtualenv creation, no `pip install` — uv handles everything.
### Running Tests
```bash
uv run pytest # run all tests
uv run pytest -v # verbose (shows each test name)
uv run pytest -m "not gpu" # skip tests that need a CUDA GPU
uv run pytest --cov # show test coverage (sources and branch mode configured in pyproject.toml)
```
Most tests run in a few seconds and don't need a GPU or model weights. Tests that require CUDA are marked with `@pytest.mark.gpu` and will be skipped automatically if no GPU is available.
### Apple Silicon (Mac) Notes
If you are contributing on an Apple Silicon Mac, there are a few extra things to be aware of.
**`uv.lock` drift:** Running `uv run pytest` on macOS regenerates `uv.lock` with macOS-specific dependency markers. **Do not commit this file.** Before staging your changes, always run:
```bash
git restore uv.lock
```
**Selecting the compute backend:** CorridorKey auto-detects MPS on Apple Silicon. To test with the MLX backend or force CPU, set the environment variable before running:
```bash
export CORRIDORKEY_BACKEND=mlx # use native MLX on Apple Silicon
export CORRIDORKEY_DEVICE=cpu # force CPU (useful for isolating device bugs)
```
**MPS operator fallback:** If PyTorch raises an error about an unsupported MPS operator, enable CPU fallback for those ops:
```bash
export PYTORCH_ENABLE_MPS_FALLBACK=1
```
### Linting and Formatting
```bash
uv run ruff check # check for lint errors
uv run ruff format --check # check formatting (no changes)
uv run ruff format # auto-format your code
```
CI runs both checks on every pull request. Running them locally before pushing saves a round-trip.
## Making Changes
### Pull Requests
1. Fork the repo and create a branch for your change
2. Make your changes
3. Run `uv run pytest` and `uv run ruff check` to make sure everything passes
4. Open a pull request against `main`
In your PR description, focus on **why** you made the change, not just what changed. If you're fixing a bug, describe the symptoms. If you're adding a feature, explain the use case. A couple of sentences is plenty.
### What Makes a Good Contribution
- **Bug fixes** — especially for edge cases in EXR/linear workflows, color space handling, or platform-specific issues
- **Tests** — more test coverage is always welcome, particularly for `clip_manager.py` and `inference_engine.py`
- **Documentation** — better explanations, usage examples, or clarifying comments in tricky code
- **Performance** — reducing GPU memory usage, speeding up frame processing, or optimizing I/O
### Code Style
- The project uses [ruff](https://docs.astral.sh/ruff/) for both linting and formatting
- Lint rules: `E, F, W, I, B` (basic style, unused imports, import sorting, common bug patterns)
- Line length: 120 characters
- Third-party code in `gvm_core/` and `VideoMaMaInferenceModule/` is excluded from lint enforcement — those are derived from research repos and we try to keep them close to upstream
### Model Weights
The model checkpoint (`CorridorKey_v1.0.pth`) and optional GVM/VideoMaMa weights are **not** in the git repo. Most tests don't need them. If you're working on inference code and need the weights, follow the download instructions in the [README](README.md).
## Questions?
Join the [Discord](https://discord.gg/zvwUrdWXJm) — it's the fastest way to get help or discuss ideas before opening a PR.
================================================
FILE: ClipsForInference/.gitkeep
================================================
================================================
FILE: CorridorKeyModule/IgnoredCheckpoints/.gitkeep
================================================
================================================
FILE: CorridorKeyModule/README.md
================================================
# CorridorKeyModule
A self-contained, high-performance AI Chroma Keying engine. This module provides a simple API to access the `CorridorKey` architecture (Hiera Backbone + CNN Refiner) for processing green screen footage.
## Features
* **Resolution Independent:** Automatically resizes input images to match the native training resolution of the model (2048x2048).
* **High Fidelity:** Preserves original input resolution using Lanczos4 resampling for final output.
* **Robust:** Supports explicit configurations for Linear (EXR) and sRGB (PNG/MP4) source inputs.
## Installation
Dependencies for the engine are managed in the main project root `requirements.txt`.
*(Requires PyTorch, NumPy, OpenCV, Timm)*
## Usage (GUI Wizard)
For most users, the easiest way to interact with the module is through the included wizard:
`clip_manager.py` (or dragging and dropping folders onto the `.bat` / `.sh` scripts).
The wizard handles finding the latest `.pth` checkpoint automatically, prompting for configuration (gamma, despill strength, despeckling), and batch processing entire sequences.
## Usage (Python API)
### 1. Initialization
Initialize the engine once. Point it to your `.pth` checkpoint. The engine is hardcoded to process at 2048x2048, representing the data it was trained on.
```python
from CorridorKeyModule import CorridorKeyEngine
# Initialize standard engine (CUDA)
engine = CorridorKeyEngine(
checkpoint_path="models/latest_model.pth",
device='cuda',
img_size=2048
)
```
### 2. Processing a Frame
The engine expects inputs as Numpy Arrays (`H, W, Channels`).
* It natively processes in **32-bit float** (`0.0 - 1.0`).
* If you pass an **8-bit integer** (`0 - 255`) array, the engine will automatically normalize it to `0.0 - 1.0` floats for you.
* If you pass a **16-bit or 32-bit float** array (like an EXR), it will process it at full precision without downgrading.
```python
import cv2
import os
# Enable EXR Support in OpenCV
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
# Load Image (Linear EXR - Read as 32-bit Float)
img_linear = cv2.imread("input.exr", cv2.IMREAD_UNCHANGED)
img_linear_rgb = cv2.cvtColor(img_linear, cv2.COLOR_BGR2RGB)
# Load Coarse Mask (Linear EXR - Read as 32-bit Float)
mask = cv2.imread("mask.exr", cv2.IMREAD_UNCHANGED)
if mask.ndim == 3:
mask = mask[:,:,0] # Keep single channel
# Process
result = engine.process_frame(
img_linear_rgb,
mask,
input_is_linear=True, # Critical: Tell the engine this is a Linear EXR
)
# Save Results (Preserving Float Precision as EXR)
# 'processed' contains the final RGBA composite (Linear 0-1 float)
proc_rgba = result['processed']
proc_bgra = cv2.cvtColor(proc_rgba, cv2.COLOR_RGBA2BGRA)
exr_flags = [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF, cv2.IMWRITE_EXR_COMPRESSION, cv2.IMWRITE_EXR_COMPRESSION_PXR24]
cv2.imwrite("output_processed.exr", proc_bgra, exr_flags)
```
## Module Structure
* `inference_engine.py`: The main API wrapper class `CorridorKeyEngine`. Handles automated input normalization (uint8 to float), tensor conversions, memory transfer, resizing to/from the 2K processing resolution, and packing the final analytical passes (RG, Alpha, Processed EXR, and Comp overlays).
* `core/model_transformer.py`: The architecture definition for the PyTorch model, combining the Hiera backbone and the convolutional refiner head.
* `core/color_utils.py`: Custom digital compositing math utilities, including logic for luminance-preserving despilling, straight/premultiplied compositing algorithms, true sRGB gamma conversions, and connected-components morphological matte cleaning.
================================================
FILE: CorridorKeyModule/__init__.py
================================================
from __future__ import annotations
from .inference_engine import CorridorKeyEngine as CorridorKeyEngine
================================================
FILE: CorridorKeyModule/backend.py
================================================
"""Backend factory — selects Torch or MLX engine and normalizes output contracts."""
from __future__ import annotations
import errno
import glob
import logging
import os
import platform
import shutil
import sys
import urllib.request
from pathlib import Path
import numpy as np
logger = logging.getLogger(__name__)
CHECKPOINT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "checkpoints")
TORCH_EXT = ".pth"
MLX_EXT = ".safetensors"
DEFAULT_IMG_SIZE = 2048
BACKEND_ENV_VAR = "CORRIDORKEY_BACKEND"
VALID_BACKENDS = ("auto", "torch", "mlx")
# Update HF_REPO_ID and HF_CHECKPOINT_FILENAME if a new model version is released.
HF_REPO_ID = "nikopueringer/CorridorKey_v1.0"
HF_CHECKPOINT_FILENAME = "CorridorKey.pth"
def resolve_backend(requested: str | None = None) -> str:
"""Resolve backend: CLI flag > env var > auto-detect.
Auto mode: Apple Silicon + corridorkey_mlx importable + .safetensors found → mlx.
Otherwise → torch.
Raises RuntimeError if explicit backend is unavailable.
"""
if requested is None or requested.lower() == "auto":
backend = os.environ.get(BACKEND_ENV_VAR, "auto").lower()
else:
backend = requested.lower()
if backend == "auto":
return _auto_detect_backend()
if backend not in VALID_BACKENDS:
raise RuntimeError(f"Unknown backend '{backend}'. Valid: {', '.join(VALID_BACKENDS)}")
if backend == "mlx":
_validate_mlx_available()
return backend
CHECKPOINT_DIR = os.path.join("CorridorKeyModule", "checkpoints")
MLX_MODEL_URL = "https://github.com/nikopueringer/corridorkey-mlx/releases/download/v1.0.0/corridorkey_mlx.safetensors"
MLX_MODEL_FILENAME = "corridorkey_mlx.safetensors"
def _auto_detect_backend() -> str:
"""Try MLX on Apple Silicon, fall back to Torch."""
if sys.platform != "darwin" or platform.machine() != "arm64":
logger.info("Not Apple Silicon — using torch backend")
return "torch"
try:
import corridorkey_mlx # type: ignore[import-not-found] # noqa: F401
except ImportError:
logger.info("corridorkey_mlx not installed — using torch backend")
return "torch"
# Auto-download logic for the .safetensors file
model_path = os.path.join(CHECKPOINT_DIR, MLX_MODEL_FILENAME)
cache_path = model_path + ".tmp"
if not os.path.exists(model_path):
logger.info(f"MLX checkpoint not found. Downloading to {model_path}...")
try:
if os.path.exists(cache_path):
os.remove(cache_path)
# Create CorridorKeyModule/checkpoints/ if it doesn't exist
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
# Download the file
urllib.request.urlretrieve(MLX_MODEL_URL, cache_path)
os.rename(cache_path, model_path)
logger.info("Download complete.")
except Exception as e:
logger.error(f"Failed to download MLX checkpoint: {e}")
logger.info("Falling back to torch backend due to download failure.")
# Clean up corrupted/partial file if the download failed midway
if os.path.exists(model_path):
os.remove(model_path)
return "torch"
logger.info("Apple Silicon + MLX available — using mlx backend")
return "mlx"
def _validate_mlx_available() -> None:
"""Raise RuntimeError with actionable message if MLX can't be used."""
if sys.platform != "darwin" or platform.machine() != "arm64":
raise RuntimeError("MLX backend requires Apple Silicon (M1+ Mac)")
try:
import corridorkey_mlx # type: ignore[import-not-found] # noqa: F401
except ImportError as err:
raise RuntimeError(
"MLX backend requested but corridorkey_mlx is not installed. "
"Install with: uv pip install corridorkey-mlx@git+https://github.com/cmoyates/corridorkey-mlx.git"
) from err
def _ensure_torch_checkpoint() -> Path:
"""Download the Torch checkpoint from HuggingFace if not present.
Returns the path to the downloaded checkpoint file.
Raises:
RuntimeError: Network or download failure.
OSError: Disk space or filesystem error.
"""
dest = Path(CHECKPOINT_DIR) / HF_CHECKPOINT_FILENAME
hf_url = f"https://huggingface.co/{HF_REPO_ID}"
from huggingface_hub import hf_hub_download
logger.info("Downloading CorridorKey checkpoint from %s ...", hf_url)
try:
cached_path = hf_hub_download(
repo_id=HF_REPO_ID,
filename=HF_CHECKPOINT_FILENAME,
)
except Exception as exc:
raise RuntimeError(
f"Failed to download CorridorKey checkpoint from {hf_url}. "
"Check your network connection and try again. "
f"Original error: {exc}"
) from exc
try:
shutil.copy2(cached_path, dest)
except OSError as exc:
if exc.errno == errno.ENOSPC:
raise OSError(
errno.ENOSPC,
"Not enough disk space to save checkpoint (~300 MB required). "
f"Free up space in {CHECKPOINT_DIR} and try again.",
) from exc
raise
logger.info("Checkpoint saved to %s", dest)
return dest
def _discover_checkpoint(ext: str) -> Path:
"""Find exactly one checkpoint with the given extension.
Raises FileNotFoundError (0 found) or ValueError (>1 found).
Includes cross-reference hints when wrong extension files exist.
"""
matches = glob.glob(os.path.join(CHECKPOINT_DIR, f"*{ext}"))
if len(matches) == 0:
if ext == TORCH_EXT:
return _ensure_torch_checkpoint()
other_ext = MLX_EXT if ext == TORCH_EXT else TORCH_EXT
other_files = glob.glob(os.path.join(CHECKPOINT_DIR, f"*{other_ext}"))
hint = ""
if other_files:
other_backend = "mlx" if other_ext == MLX_EXT else "torch"
hint = f" (Found {other_ext} files — did you mean --backend={other_backend}?)"
raise FileNotFoundError(f"No {ext} checkpoint found in {CHECKPOINT_DIR}.{hint}")
if len(matches) > 1:
names = [os.path.basename(f) for f in matches]
raise ValueError(f"Multiple {ext} checkpoints in {CHECKPOINT_DIR}: {names}. Keep exactly one.")
return Path(matches[0])
def _wrap_mlx_output(raw: dict, despill_strength: float, auto_despeckle: bool, despeckle_size: int) -> dict:
"""Normalize MLX uint8 output to match Torch float32 contract.
Torch contract:
alpha: [H,W,1] float32 0-1
fg: [H,W,3] float32 0-1 sRGB
comp: [H,W,3] float32 0-1 sRGB
processed: [H,W,4] float32 linear premul RGBA
"""
from CorridorKeyModule.core import color_utils as cu
# alpha: uint8 [H,W] → float32 [H,W,1]
alpha_raw = raw["alpha"]
alpha = alpha_raw.astype(np.float32) / 255.0
if alpha.ndim == 2:
alpha = alpha[:, :, np.newaxis]
# fg: uint8 [H,W,3] → float32 [H,W,3] (sRGB)
fg = raw["fg"].astype(np.float32) / 255.0
# Apply despeckle (MLX stubs this)
if auto_despeckle:
processed_alpha = cu.clean_matte(alpha, area_threshold=despeckle_size, dilation=25, blur_size=5)
else:
processed_alpha = alpha
# Apply despill (MLX stubs this)
fg_despilled = cu.despill(fg, green_limit_mode="average", strength=despill_strength)
# Composite over checkerboard for comp output
h, w = fg.shape[:2]
bg_srgb = cu.create_checkerboard(w, h, checker_size=128, color1=0.15, color2=0.55)
bg_lin = cu.srgb_to_linear(bg_srgb)
fg_despilled_lin = cu.srgb_to_linear(fg_despilled)
comp_lin = cu.composite_straight(fg_despilled_lin, bg_lin, processed_alpha)
comp_srgb = cu.linear_to_srgb(comp_lin)
# Build processed: [H,W,4] linear premul RGBA
fg_premul_lin = cu.premultiply(fg_despilled_lin, processed_alpha)
processed_rgba = np.concatenate([fg_premul_lin, processed_alpha], axis=-1)
return {
"alpha": alpha, # raw prediction (before despeckle), matches Torch
"fg": fg, # raw sRGB prediction, matches Torch
"comp": comp_srgb, # sRGB composite on checker
"processed": processed_rgba, # linear premul RGBA
}
class _MLXEngineAdapter:
"""Wraps CorridorKeyMLXEngine to match Torch output contract."""
def __init__(self, raw_engine):
self._engine = raw_engine
logger.info("MLX adapter active: despill and despeckle are handled by the adapter layer, not native MLX")
def process_frame(
self,
image,
mask_linear,
refiner_scale=1.0,
input_is_linear=False,
fg_is_straight=True,
despill_strength=1.0,
auto_despeckle=True,
despeckle_size=400,
):
"""Delegate to MLX engine, then normalize output to Torch contract."""
# MLX engine expects uint8 input — convert if float
if image.dtype != np.uint8:
image_u8 = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8)
else:
image_u8 = image
if mask_linear.dtype != np.uint8:
mask_u8 = (np.clip(mask_linear, 0.0, 1.0) * 255).astype(np.uint8)
else:
mask_u8 = mask_linear
# Squeeze mask to 2D for MLX (it validates [H,W] or [H,W,1])
if mask_u8.ndim == 3:
mask_u8 = mask_u8[:, :, 0]
raw = self._engine.process_frame(
image_u8,
mask_u8,
refiner_scale=refiner_scale,
input_is_linear=input_is_linear,
fg_is_straight=fg_is_straight,
despill_strength=0.0, # disable MLX stubs — adapter applies these
auto_despeckle=False,
despeckle_size=despeckle_size,
)
return _wrap_mlx_output(raw, despill_strength, auto_despeckle, despeckle_size)
DEFAULT_MLX_TILE_SIZE = 512
DEFAULT_MLX_TILE_OVERLAP = 64
def create_engine(
backend: str | None = None,
device: str | None = None,
img_size: int = DEFAULT_IMG_SIZE,
tile_size: int | None = DEFAULT_MLX_TILE_SIZE,
overlap: int = DEFAULT_MLX_TILE_OVERLAP,
):
"""Factory: returns an engine with process_frame() matching the Torch contract.
Args:
tile_size: MLX only — tile size for tiled inference (default 512).
Set to None to disable tiling and use full-frame inference.
overlap: MLX only — overlap pixels between tiles (default 64).
"""
backend = resolve_backend(backend)
if backend == "mlx":
ckpt = _discover_checkpoint(MLX_EXT)
from corridorkey_mlx import CorridorKeyMLXEngine # type: ignore[import-not-found]
raw_engine = CorridorKeyMLXEngine(str(ckpt), img_size=img_size, tile_size=tile_size, overlap=overlap)
mode = f"tiled (tile={tile_size}, overlap={overlap})" if tile_size else "full-frame"
logger.info("MLX engine loaded: %s [%s]", ckpt.name, mode)
return _MLXEngineAdapter(raw_engine)
else:
ckpt = _discover_checkpoint(TORCH_EXT)
from CorridorKeyModule.inference_engine import CorridorKeyEngine
logger.info("Torch engine loaded: %s (device=%s)", ckpt.name, device)
return CorridorKeyEngine(checkpoint_path=str(ckpt), device=device or "cpu", img_size=img_size)
================================================
FILE: CorridorKeyModule/checkpoints/.gitkeep
================================================
================================================
FILE: CorridorKeyModule/core/__init__.py
================================================
================================================
FILE: CorridorKeyModule/core/color_utils.py
================================================
from __future__ import annotations
import functools
from collections.abc import Callable
import cv2
import numpy as np
import torch
def _is_tensor(x: np.ndarray | torch.Tensor) -> bool:
return isinstance(x, torch.Tensor)
def _if_tensor(is_tensor: bool, tensor_func: Callable, numpy_func: Callable) -> Callable:
return tensor_func if is_tensor else numpy_func
def _power(x: np.ndarray | torch.Tensor, exponent: float) -> np.ndarray | torch.Tensor:
"""
Power function that supports both Numpy arrays and PyTorch tensors.
"""
power = _if_tensor(_is_tensor(x), torch.pow, np.power)
return power(x, exponent)
def _where(
condition: np.ndarray | torch.Tensor, x: np.ndarray | torch.Tensor, y: np.ndarray | torch.Tensor
) -> np.ndarray | torch.Tensor:
"""
Where function that supports both Numpy arrays and PyTorch tensors.
"""
where = _if_tensor(_is_tensor(x), torch.where, np.where)
return where(condition, x, y)
def _clamp(x: np.ndarray | torch.Tensor, min: float) -> np.ndarray | torch.Tensor:
"""
Clamp function that supports both Numpy arrays and PyTorch tensors.
"""
if isinstance(x, torch.Tensor):
return x.clamp(min=0.0)
return np.clip(x, 0.0, None)
_torch_stack = functools.partial(torch.stack, dim=-1)
_numpy_stack = functools.partial(np.stack, axis=-1)
def linear_to_srgb(x: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
"""
Converts Linear to sRGB using the official piecewise sRGB transfer function.
Supports both Numpy arrays and PyTorch tensors.
"""
x = _clamp(x, 0.0)
mask = x <= 0.0031308
return _where(mask, x * 12.92, 1.055 * _power(x, 1.0 / 2.4) - 0.055)
def srgb_to_linear(x: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
"""
Converts sRGB to Linear using the official piecewise sRGB transfer function.
Supports both Numpy arrays and PyTorch tensors.
"""
x = _clamp(x, 0.0)
mask = x <= 0.04045
return _where(mask, x / 12.92, _power((x + 0.055) / 1.055, 2.4))
def premultiply(fg: np.ndarray | torch.Tensor, alpha: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
"""
Premultiplies foreground by alpha.
fg: Color [..., C] or [C, ...]
alpha: Alpha [..., 1] or [1, ...]
"""
return fg * alpha
def unpremultiply(
fg: np.ndarray | torch.Tensor, alpha: np.ndarray | torch.Tensor, eps: float = 1e-6
) -> np.ndarray | torch.Tensor:
"""
Un-premultiplies foreground by alpha.
Ref: fg_straight = fg_premul / (alpha + eps)
"""
return fg / (alpha + eps)
def composite_straight(
fg: np.ndarray | torch.Tensor, bg: np.ndarray | torch.Tensor, alpha: np.ndarray | torch.Tensor
) -> np.ndarray | torch.Tensor:
"""
Composites Straight FG over BG.
Formula: FG * Alpha + BG * (1 - Alpha)
"""
return fg * alpha + bg * (1.0 - alpha)
def composite_premul(
fg: np.ndarray | torch.Tensor, bg: np.ndarray | torch.Tensor, alpha: np.ndarray | torch.Tensor
) -> np.ndarray | torch.Tensor:
"""
Composites Premultiplied FG over BG.
Formula: FG + BG * (1 - Alpha)
"""
return fg + bg * (1.0 - alpha)
def rgb_to_yuv(image: torch.Tensor) -> torch.Tensor:
"""
Converts RGB to YUV (Rec. 601).
Input: [..., 3, H, W] or [..., 3] depending on layout.
Supports standard PyTorch BCHW.
"""
if not _is_tensor(image):
raise TypeError("rgb_to_yuv only supports dict/tensor inputs currently")
# Weights for RGB -> Y
# Rec. 601: 0.299, 0.587, 0.114
# Assume BCHW layout if 4 dims
if image.dim() == 4:
r = image[:, 0:1, :, :]
g = image[:, 1:2, :, :]
b = image[:, 2:3, :, :]
elif image.dim() == 3 and image.shape[0] == 3: # CHW
r = image[0:1, :, :]
g = image[1:2, :, :]
b = image[2:3, :, :]
else:
# Last dim conversion
r = image[..., 0]
g = image[..., 1]
b = image[..., 2]
y = 0.299 * r + 0.587 * g + 0.114 * b
u = 0.492 * (b - y)
v = 0.877 * (r - y)
if image.dim() >= 3 and image.shape[-3] == 3: # Concatenate along Channel dim
return torch.cat([y, u, v], dim=-3)
else:
return torch.stack([y, u, v], dim=-1)
def dilate_mask(mask: np.ndarray | torch.Tensor, radius: int) -> np.ndarray | torch.Tensor:
"""
Dilates a mask by a given radius.
Supports Numpy (using cv2) and PyTorch (using MaxPool).
radius: Int (pixels). 0 = No change.
"""
if radius <= 0:
return mask
kernel_size = int(radius * 2 + 1)
if isinstance(mask, torch.Tensor):
# PyTorch Dilation (using Max Pooling)
# Expects [B, C, H, W]
orig_dim = mask.dim()
if orig_dim == 2:
mask = mask.unsqueeze(0).unsqueeze(0)
elif orig_dim == 3:
mask = mask.unsqueeze(0)
padding = radius
dilated = torch.nn.functional.max_pool2d(mask, kernel_size, stride=1, padding=padding)
if orig_dim == 2:
return dilated.squeeze()
elif orig_dim == 3:
return dilated.squeeze(0)
return dilated
# Numpy Dilation (using OpenCV)
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
return cv2.dilate(mask, kernel)
def apply_garbage_matte(
predicted_matte: np.ndarray | torch.Tensor,
garbage_matte_input: np.ndarray | torch.Tensor | None,
dilation: int = 10,
) -> np.ndarray | torch.Tensor:
"""
Multiplies predicted matte by a dilated garbage matte to clean up background.
"""
if garbage_matte_input is None:
return predicted_matte
garbage_mask = dilate_mask(garbage_matte_input, dilation)
# Ensure dimensions match for multiplication
if _is_tensor(predicted_matte):
# Handle broadcasting if needed
pass
elif garbage_mask.ndim == 2 and predicted_matte.ndim == 3:
# Numpy
garbage_mask = garbage_mask[:, :, np.newaxis]
return predicted_matte * garbage_mask
def despill(
image: np.ndarray | torch.Tensor, green_limit_mode: str = "average", strength: float = 1.0
) -> np.ndarray | torch.Tensor:
"""
Removes green spill from an RGB image using a luminance-preserving method.
image: RGB float (0-1).
green_limit_mode: 'average' ((R+B)/2) or 'max' (max(R, B)).
strength: 0.0 to 1.0 multiplier for the despill effect.
"""
if strength <= 0.0:
return image
tensor = _is_tensor(image)
_maximum = _if_tensor(tensor, torch.max, np.maximum)
_stack = _if_tensor(tensor, _torch_stack, _numpy_stack)
r = image[..., 0]
g = image[..., 1]
b = image[..., 2]
if green_limit_mode == "max":
limit = _maximum(r, b)
else:
limit = (r + b) / 2.0
if isinstance(image, torch.Tensor):
# PyTorch Impl — g/limit are Tensor since image is Tensor
diff: torch.Tensor = g - limit # type: ignore[assignment]
spill_amount = torch.clamp(diff, min=0.0)
else:
# Numpy Impl
spill_amount = np.maximum(g - limit, 0.0)
g_new = g - spill_amount
r_new = r + (spill_amount * 0.5)
b_new = b + (spill_amount * 0.5)
despilled = _stack([r_new, g_new, b_new])
if strength < 1.0:
return image * (1.0 - strength) + despilled * strength
return despilled
def clean_matte(alpha_np: np.ndarray, area_threshold: int = 300, dilation: int = 15, blur_size: int = 5) -> np.ndarray:
"""
Cleans up small disconnected components (like tracking markers) from a predicted alpha matte.
alpha_np: Numpy array [H, W] or [H, W, 1] float (0.0 - 1.0)
"""
# Needs to be 2D
is_3d = False
if alpha_np.ndim == 3:
is_3d = True
alpha_np = alpha_np[:, :, 0]
# Threshold to binary
mask_8u = (alpha_np > 0.5).astype(np.uint8) * 255
# Find connected components
num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask_8u, connectivity=8)
# Create an empty mask for the cleaned components
cleaned_mask = np.zeros_like(mask_8u)
# Keep components larger than the threshold (skip label 0, which is background)
for i in range(1, num_labels):
if stats[i, cv2.CC_STAT_AREA] >= area_threshold:
cleaned_mask[labels == i] = 255
# Dilate
if dilation > 0:
kernel_size = int(dilation * 2 + 1)
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
cleaned_mask = cv2.dilate(cleaned_mask, kernel)
# Blur
if blur_size > 0:
b_size = int(blur_size * 2 + 1)
cleaned_mask = cv2.GaussianBlur(cleaned_mask, (b_size, b_size), 0)
# Convert back to 0-1 float
safe_zone = cleaned_mask.astype(np.float32) / 255.0
# Multiply original alpha by the safe zone
result_alpha = alpha_np * safe_zone
if is_3d:
result_alpha = result_alpha[:, :, np.newaxis]
return result_alpha
def create_checkerboard(
width: int, height: int, checker_size: int = 64, color1: float = 0.2, color2: float = 0.4
) -> np.ndarray:
"""
Creates a linear grayscale checkerboard pattern.
Returns: Numpy array [H, W, 3] float (0.0-1.0)
"""
# Create coordinate grids
x = np.arange(width)
y = np.arange(height)
# Determine tile parity
x_tiles = x // checker_size
y_tiles = y // checker_size
# Broadcast to 2D
x_grid, y_grid = np.meshgrid(x_tiles, y_tiles)
# XOR for checker pattern (1 if odd, 0 if even)
checker = (x_grid + y_grid) % 2
# Map 0 to color1 and 1 to color2
bg_img = np.where(checker == 0, color1, color2).astype(np.float32)
# Make it 3-channel
return np.stack([bg_img, bg_img, bg_img], axis=-1)
================================================
FILE: CorridorKeyModule/core/model_transformer.py
================================================
from __future__ import annotations
import logging
import timm
import torch
import torch.nn as nn
import torch.nn.functional as F
logger = logging.getLogger(__name__)
class MLP(nn.Module):
"""Linear embedding: C_in -> C_out."""
def __init__(self, input_dim: int = 2048, embed_dim: int = 768) -> None:
super().__init__()
self.proj = nn.Linear(input_dim, embed_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.proj(x)
class DecoderHead(nn.Module):
def __init__(
self, feature_channels: list[int] | None = None, embedding_dim: int = 256, output_dim: int = 1
) -> None:
super().__init__()
if feature_channels is None:
feature_channels = [112, 224, 448, 896]
# MLP layers to unify channel dimensions
self.linear_c4 = MLP(input_dim=feature_channels[3], embed_dim=embedding_dim)
self.linear_c3 = MLP(input_dim=feature_channels[2], embed_dim=embedding_dim)
self.linear_c2 = MLP(input_dim=feature_channels[1], embed_dim=embedding_dim)
self.linear_c1 = MLP(input_dim=feature_channels[0], embed_dim=embedding_dim)
# Fuse
self.linear_fuse = nn.Conv2d(embedding_dim * 4, embedding_dim, kernel_size=1, bias=False)
self.bn = nn.BatchNorm2d(embedding_dim)
self.relu = nn.ReLU(inplace=True)
# Predict
self.dropout = nn.Dropout(0.1)
self.classifier = nn.Conv2d(embedding_dim, output_dim, kernel_size=1)
def forward(self, features: list[torch.Tensor]) -> torch.Tensor:
c1, c2, c3, c4 = features
n, _, h, w = c4.shape
# Resize to C1 size (which is H/4)
_c4 = self.linear_c4(c4.flatten(2).transpose(1, 2)).transpose(1, 2).view(n, -1, c4.shape[2], c4.shape[3])
_c4 = F.interpolate(_c4, size=c1.shape[2:], mode="bilinear", align_corners=False)
_c3 = self.linear_c3(c3.flatten(2).transpose(1, 2)).transpose(1, 2).view(n, -1, c3.shape[2], c3.shape[3])
_c3 = F.interpolate(_c3, size=c1.shape[2:], mode="bilinear", align_corners=False)
_c2 = self.linear_c2(c2.flatten(2).transpose(1, 2)).transpose(1, 2).view(n, -1, c2.shape[2], c2.shape[3])
_c2 = F.interpolate(_c2, size=c1.shape[2:], mode="bilinear", align_corners=False)
_c1 = self.linear_c1(c1.flatten(2).transpose(1, 2)).transpose(1, 2).view(n, -1, c1.shape[2], c1.shape[3])
_c = self.linear_fuse(torch.cat([_c4, _c3, _c2, _c1], dim=1))
_c = self.bn(_c)
_c = self.relu(_c)
x = self.dropout(_c)
x = self.classifier(x)
return x
class RefinerBlock(nn.Module):
"""
Residual Block with Dilation and GroupNorm (Safe for Batch Size 2).
"""
def __init__(self, channels: int, dilation: int = 1) -> None:
super().__init__()
self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=dilation, dilation=dilation)
self.gn1 = nn.GroupNorm(8, channels)
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=dilation, dilation=dilation)
self.gn2 = nn.GroupNorm(8, channels)
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x
out = self.conv1(x)
out = self.gn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.gn2(out)
out += residual
out = self.relu(out)
return out
class CNNRefinerModule(nn.Module):
"""
Dilated Residual Refiner (Receptive Field ~65px).
designed to solve Macroblocking artifacts from Hiera.
Structure: Stem -> Res(d1) -> Res(d2) -> Res(d4) -> Res(d8) -> Projection.
"""
def __init__(self, in_channels: int = 7, hidden_channels: int = 64, out_channels: int = 4) -> None:
super().__init__()
# Stem
self.stem = nn.Sequential(
nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1),
nn.GroupNorm(8, hidden_channels),
nn.ReLU(inplace=True),
)
# Dilated Residual Blocks (RF Expansion)
self.res1 = RefinerBlock(hidden_channels, dilation=1)
self.res2 = RefinerBlock(hidden_channels, dilation=2)
self.res3 = RefinerBlock(hidden_channels, dilation=4)
self.res4 = RefinerBlock(hidden_channels, dilation=8)
# Final Projection (No Activation, purely additive logits)
self.final = nn.Conv2d(hidden_channels, out_channels, kernel_size=1)
# Tiny Noise Init (Whisper) - Provides gradients without shock
nn.init.normal_(self.final.weight, mean=0.0, std=1e-3)
nn.init.constant_(self.final.bias, 0)
def forward(self, img: torch.Tensor, coarse_pred: torch.Tensor) -> torch.Tensor:
# img: [B, 3, H, W]
# coarse_pred: [B, 4, H, W]
x = torch.cat([img, coarse_pred], dim=1)
x = self.stem(x)
x = self.res1(x)
x = self.res2(x)
x = self.res3(x)
x = self.res4(x)
# Output Scaling (10x Boost)
# Allows the Refiner to predict small stable values (e.g. 0.5) that become strong corrections (5.0).
return self.final(x) * 10.0
class GreenFormer(nn.Module):
def __init__(
self,
encoder_name: str = "hiera_base_plus_224.mae_in1k_ft_in1k",
in_channels: int = 4,
img_size: int = 512,
use_refiner: bool = True,
) -> None:
super().__init__()
# --- Encoder ---
# Load Pretrained Hiera
# 1. Create Target Model (512x512, Random Weights)
# We use features_only=True, which wraps it in FeatureGetterNet
logger.info("Initializing %s (img_size=%d)", encoder_name, img_size)
self.encoder = timm.create_model(encoder_name, pretrained=False, features_only=True, img_size=img_size)
# We skip downloading/loading base weights because the user's checkpoint
# (loaded immediately after this) contains all weights, including correctly
# trained/sized PosEmbeds. This keeps the project offline-capable using only local assets.
logger.info("Skipped downloading base weights (relying on custom checkpoint)")
# Patch First Layer for 4 channels
if in_channels != 3:
self._patch_input_layer(in_channels)
# Get feature info
# Verified Hiera Base Plus channels: [112, 224, 448, 896]
# We can try to fetch dynamically
try:
feature_channels = self.encoder.feature_info.channels()
except (AttributeError, TypeError):
feature_channels = [112, 224, 448, 896]
logger.info("Feature channels: %s", feature_channels)
# --- Decoders ---
embedding_dim = 256
# Alpha Decoder (Outputs 1 channel)
self.alpha_decoder = DecoderHead(feature_channels, embedding_dim, output_dim=1)
# Foreground Decoder (Outputs 3 channels)
self.fg_decoder = DecoderHead(feature_channels, embedding_dim, output_dim=3)
# --- Refiner ---
# CNN Refiner
# In Channels: 3 (RGB) + 4 (Coarse Pred) = 7
self.use_refiner = use_refiner
if self.use_refiner:
self.refiner = CNNRefinerModule(in_channels=7, hidden_channels=64, out_channels=4)
else:
self.refiner = None
logger.info("Refiner module DISABLED (backbone-only mode)")
def _patch_input_layer(self, in_channels: int) -> None:
"""
Modifies the first convolution layer to accept `in_channels`.
Copies existing RGB weights and initializes extras to zero.
"""
# Hiera: self.encoder.model.patch_embed.proj
try:
patch_embed = self.encoder.model.patch_embed.proj
except AttributeError:
# Fallback if timm changes structure or for other models
patch_embed = self.encoder.patch_embed.proj
weight = patch_embed.weight.data # [Out, 3, K, K]
bias = patch_embed.bias.data if patch_embed.bias is not None else None
new_in_channels = in_channels
out_channels, _, k, k = weight.shape
# Create new conv
new_conv = nn.Conv2d(
new_in_channels,
out_channels,
kernel_size=k,
stride=patch_embed.stride,
padding=patch_embed.padding,
bias=(bias is not None),
)
# Copy weights
new_conv.weight.data[:, :3, :, :] = weight
# Initialize new channels to 0 (Weight Patching)
new_conv.weight.data[:, 3:, :, :] = 0.0
if bias is not None:
new_conv.bias.data = bias
# Replace in module
try:
self.encoder.model.patch_embed.proj = new_conv
except AttributeError:
self.encoder.patch_embed.proj = new_conv
logger.info("Patched input layer: 3 → %d channels (extra initialized to 0)", in_channels)
def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
# x: [B, 4, H, W]
input_size = x.shape[2:]
# Encode
features = self.encoder(x) # Returns list of features
# Decode Streams
alpha_logits = self.alpha_decoder(features) # [B, 1, H/4, W/4]
fg_logits = self.fg_decoder(features) # [B, 3, H/4, W/4]
# Upsample to full resolution (Bilinear)
# These are the "Coarse" LOGITS
alpha_logits_up = F.interpolate(alpha_logits, size=input_size, mode="bilinear", align_corners=False)
fg_logits_up = F.interpolate(fg_logits, size=input_size, mode="bilinear", align_corners=False)
# --- HUMILITY CLAMP REMOVED (Phase 3) ---
# User requested NO CLAMPING to preserve all backbone detail.
# Refiner sees raw logits (-inf to +inf).
# alpha_logits_up = torch.clamp(alpha_logits_up, -3.0, 3.0)
# fg_logits_up = torch.clamp(fg_logits_up, -3.0, 3.0)
# Coarse Probs (for Loss and Refiner Input)
alpha_coarse = torch.sigmoid(alpha_logits_up)
fg_coarse = torch.sigmoid(fg_logits_up)
# --- Refinement (CNN Hybrid) ---
# 4. Refine (CNN)
# Input to refiner: RGB Image (first 3 channels of x) + Coarse Predictions (Probs)
# We give the refiner 'Probs' as input features because they are normalized [0,1]
rgb = x[:, :3, :, :]
# Feed the Refiner
coarse_pred = torch.cat([alpha_coarse, fg_coarse], dim=1) # [B, 4, H, W]
# Refiner outputs DELTA LOGITS
# The refiner predicts the correction in valid score space (-inf, inf)
if self.use_refiner and self.refiner is not None:
delta_logits = self.refiner(rgb, coarse_pred)
else:
# Zero Deltas
delta_logits = torch.zeros_like(coarse_pred)
delta_alpha = delta_logits[:, 0:1]
delta_fg = delta_logits[:, 1:4]
# Residual Addition in Logit Space
# This allows infinite correction capability and prevents saturation blocking
alpha_final_logits = alpha_logits_up + delta_alpha
fg_final_logits = fg_logits_up + delta_fg
# Final Activation
alpha_final = torch.sigmoid(alpha_final_logits)
fg_final = torch.sigmoid(fg_final_logits)
return {"alpha": alpha_final, "fg": fg_final}
================================================
FILE: CorridorKeyModule/inference_engine.py
================================================
from __future__ import annotations
import logging
import math
import os
import sys
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from .core import color_utils as cu
from .core.model_transformer import GreenFormer
logger = logging.getLogger(__name__)
class CorridorKeyEngine:
def __init__(
self,
checkpoint_path: str,
device: str = "cpu",
img_size: int = 2048,
use_refiner: bool = True,
mixed_precision: bool = True,
model_precision: torch.dtype = torch.float32,
) -> None:
self.device = torch.device(device)
self.img_size = img_size
self.checkpoint_path = checkpoint_path
self.use_refiner = use_refiner
self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)
if mixed_precision or model_precision != torch.float32:
# Use faster matrix multiplication implementation
# This reduces the floating point precision a little bit,
# but it should be negligible compared to fp16 precision
torch.set_float32_matmul_precision("high")
self.mixed_precision = mixed_precision
if mixed_precision and model_precision == torch.float16:
# using mixed precision, when the precision is already fp16, is slower
self.mixed_precision = False
self.model_precision = model_precision
model = self._load_model().to(model_precision)
# We only tested compilation on windows and linux. For other platforms compilation is disabled as a precaution.
if sys.platform == "linux" or sys.platform == "win32":
# Try compiling the model. Fallback to eager mode if it fails.
try:
self.model = torch.compile(model)
# Trigger compilation with a dummy input
dummy_input = torch.zeros(1, 4, img_size, img_size, dtype=model_precision, device=self.device)
with torch.inference_mode():
self.model(dummy_input)
except Exception as e:
logger.info(f"Model compilation failed with error: {e}")
logger.warning("Model compilation failed. Falling back to eager mode.")
torch.cuda.empty_cache()
self.model = model
def _load_model(self) -> GreenFormer:
logger.info("Loading CorridorKey from %s", self.checkpoint_path)
# Initialize Model (Hiera Backbone)
model = GreenFormer(
encoder_name="hiera_base_plus_224.mae_in1k_ft_in1k", img_size=self.img_size, use_refiner=self.use_refiner
)
model = model.to(self.device)
model.eval()
# Load Weights
if not os.path.isfile(self.checkpoint_path):
raise FileNotFoundError(f"Checkpoint not found: {self.checkpoint_path}")
checkpoint = torch.load(self.checkpoint_path, map_location=self.device, weights_only=True)
state_dict = checkpoint.get("state_dict", checkpoint)
# Fix Compiled Model Prefix & Handle PosEmbed Mismatch
new_state_dict = {}
model_state = model.state_dict()
for k, v in state_dict.items():
if k.startswith("_orig_mod."):
k = k[10:]
# Check for PosEmbed Mismatch
if "pos_embed" in k and k in model_state:
if v.shape != model_state[k].shape:
print(f"Resizing {k} from {v.shape} to {model_state[k].shape}")
# v: [1, N_src, C]
# target: [1, N_dst, C]
# We assume square grid
N_src = v.shape[1]
N_dst = model_state[k].shape[1]
C = v.shape[2]
grid_src = int(math.sqrt(N_src))
grid_dst = int(math.sqrt(N_dst))
# Reshape to [1, C, H, W]
v_img = v.permute(0, 2, 1).view(1, C, grid_src, grid_src)
# Interpolate
v_resized = F.interpolate(v_img, size=(grid_dst, grid_dst), mode="bicubic", align_corners=False)
# Reshape back
v = v_resized.flatten(2).transpose(1, 2)
new_state_dict[k] = v
missing, unexpected = model.load_state_dict(new_state_dict, strict=False)
if len(missing) > 0:
print(f"[Warning] Missing keys: {missing}")
if len(unexpected) > 0:
print(f"[Warning] Unexpected keys: {unexpected}")
return model
@torch.inference_mode()
def process_frame(
self,
image: np.ndarray,
mask_linear: np.ndarray,
refiner_scale: float = 1.0,
input_is_linear: bool = False,
fg_is_straight: bool = True,
despill_strength: float = 1.0,
auto_despeckle: bool = True,
despeckle_size: int = 400,
) -> dict[str, np.ndarray]:
"""
Process a single frame.
Args:
image: Numpy array [H, W, 3] (0.0-1.0 or 0-255).
- If input_is_linear=False (Default): Assumed sRGB.
- If input_is_linear=True: Assumed Linear.
mask_linear: Numpy array [H, W] or [H, W, 1] (0.0-1.0). Assumed Linear.
refiner_scale: Multiplier for Refiner Deltas (default 1.0).
input_is_linear: bool. If True, resizes in Linear then transforms to sRGB.
If False, resizes in sRGB (standard).
fg_is_straight: bool. If True, assumes FG output is Straight (unpremultiplied).
If False, assumes FG output is Premultiplied.
despill_strength: float. 0.0 to 1.0 multiplier for the despill effect.
auto_despeckle: bool. If True, cleans up small disconnected components from the predicted alpha matte.
despeckle_size: int. Minimum number of consecutive pixels required to keep an island.
Returns:
dict: {'alpha': np, 'fg': np (sRGB), 'comp': np (sRGB on Gray)}
"""
# 1. Inputs Check & Normalization
if image.dtype == np.uint8:
image = image.astype(np.float32) / 255.0
if mask_linear.dtype == np.uint8:
mask_linear = mask_linear.astype(np.float32) / 255.0
h, w = image.shape[:2]
# Ensure Mask Shape
if mask_linear.ndim == 2:
mask_linear = mask_linear[:, :, np.newaxis]
# 2. Resize to Model Size
# If input is linear, we resize in linear to preserve energy/highlights,
# THEN convert to sRGB for the model.
if input_is_linear:
# Resize in Linear
img_resized_lin = cv2.resize(image, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR)
# Convert to sRGB for Model
img_resized = cu.linear_to_srgb(img_resized_lin)
else:
# Standard sRGB Resize
img_resized = cv2.resize(image, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR)
mask_resized = cv2.resize(mask_linear, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR)
if mask_resized.ndim == 2:
mask_resized = mask_resized[:, :, np.newaxis]
# 3. Normalize (ImageNet)
# Model expects sRGB input normalized
img_norm = (img_resized - self.mean) / self.std
# 4. Prepare Tensor
inp_np = np.concatenate([img_norm, mask_resized], axis=-1) # [H, W, 4]
inp_t = torch.from_numpy(inp_np.transpose((2, 0, 1))).unsqueeze(0).to(self.model_precision).to(self.device)
# 5. Inference
# Hook for Refiner Scaling
handle = None
if refiner_scale != 1.0 and self.model.refiner is not None:
def scale_hook(module, input, output):
return output * refiner_scale
handle = self.model.refiner.register_forward_hook(scale_hook)
with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=self.mixed_precision):
out = self.model(inp_t)
if handle:
handle.remove()
pred_alpha = out["alpha"]
pred_fg = out["fg"] # Output is sRGB (Sigmoid)
# 6. Post-Process (Resize Back to Original Resolution)
# We use Lanczos4 for high-quality resampling to minimize blur when going back to 4K/Original.
res_alpha = pred_alpha[0].permute(1, 2, 0).float().cpu().numpy()
res_fg = pred_fg[0].permute(1, 2, 0).float().cpu().numpy()
res_alpha = cv2.resize(res_alpha, (w, h), interpolation=cv2.INTER_LANCZOS4)
res_fg = cv2.resize(res_fg, (w, h), interpolation=cv2.INTER_LANCZOS4)
if res_alpha.ndim == 2:
res_alpha = res_alpha[:, :, np.newaxis]
# --- ADVANCED COMPOSITING ---
# A. Clean Matte (Auto-Despeckle)
if auto_despeckle:
processed_alpha = cu.clean_matte(res_alpha, area_threshold=despeckle_size, dilation=25, blur_size=5)
else:
processed_alpha = res_alpha
# B. Despill FG
# res_fg is sRGB.
fg_despilled = cu.despill(res_fg, green_limit_mode="average", strength=despill_strength)
# C. Premultiply (for EXR Output)
# CONVERT TO LINEAR FIRST! EXRs must house linear color premultiplied by linear alpha.
fg_despilled_lin = cu.srgb_to_linear(fg_despilled)
fg_premul_lin = cu.premultiply(fg_despilled_lin, processed_alpha)
# D. Pack RGBA
# [H, W, 4] - All channels are now strictly Linear Float
processed_rgba = np.concatenate([fg_premul_lin, processed_alpha], axis=-1)
# ----------------------------
# 7. Composite (on Checkerboard) for checking
# Generate Dark/Light Gray Checkerboard (in sRGB, convert to Linear)
bg_srgb = cu.create_checkerboard(w, h, checker_size=128, color1=0.15, color2=0.55)
bg_lin = cu.srgb_to_linear(bg_srgb)
if fg_is_straight:
comp_lin = cu.composite_straight(fg_despilled_lin, bg_lin, processed_alpha)
else:
# If premultiplied model, we shouldn't multiply again (though our pipeline forces straight)
comp_lin = cu.composite_premul(fg_despilled_lin, bg_lin, processed_alpha)
comp_srgb = cu.linear_to_srgb(comp_lin)
return { # type: ignore[return-value] # cu.* returns ndarray|Tensor but inputs are always ndarray here
"alpha": res_alpha, # Linear, Raw Prediction
"fg": res_fg, # sRGB, Raw Prediction (Straight)
"comp": comp_srgb, # sRGB, Composite
"processed": processed_rgba, # Linear/Premul, RGBA, Garbage Matted & Despilled
}
================================================
FILE: CorridorKey_DRAG_CLIPS_HERE_local.bat
================================================
@echo off
REM Corridor Key Launcher - Local
REM Set script path (assumes corridorkey_cli.py is in the same directory as this batch file)
set "SCRIPT_DIR=%~dp0"
set "LOCAL_SCRIPT=%SCRIPT_DIR%corridorkey_cli.py"
REM SAFETY CHECK: Ensure a folder was dragged onto the script
if "%~1"=="" (
echo [ERROR] No target folder provided.
echo.
echo USAGE:
echo Please DRAG AND DROP a folder onto this script to process it.
echo Do not double-click this script directly.
echo.
pause
exit /b
)
REM Folder dragged? Use it as the target path.
set "WIN_PATH=%~1"
REM Strip trailing slash if present
if "%WIN_PATH:~-1%"=="\" set "WIN_PATH=%WIN_PATH:~0,-1%"
echo Starting Corridor Key locally...
echo Target: "%WIN_PATH%"
REM Run via uv entry point (handles the virtual environment automatically)
cd /d "%SCRIPT_DIR%"
uv run --extra cuda corridorkey wizard "%WIN_PATH%"
pause
================================================
FILE: CorridorKey_DRAG_CLIPS_HERE_local.sh
================================================
#!/usr/bin/env bash
# Corridor Key Launcher - Local Linux/macOS
cd "$(dirname "$0")"
# Get the directory where this script is located
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
LOCAL_SCRIPT="$SCRIPT_DIR/corridorkey_cli.py"
# SAFETY CHECK: Ensure a folder was provided as an argument
if [ -z "$1" ]; then
echo "[ERROR] No target folder provided."
echo ""
echo "USAGE:"
echo "You can either run this script from the terminal and provide a path:"
echo " ./CorridorKey_DRAG_CLIPS_HERE_local.sh /path/to/your/clip/folder"
echo ""
echo "Or, in many Linux/macOS desktop environments, you can simply"
echo "DRAG AND DROP a folder onto this script icon to process it."
echo ""
read -p "Press enter to exit..."
exit 1
fi
# Folder dragged or provided via CLI? Use it as the target path.
TARGET_PATH="$1"
# Strip trailing slash if present
TARGET_PATH="${TARGET_PATH%/}"
# Ensure uv is available before attempting to run
if ! command -v uv &> /dev/null; then
echo "[ERROR] 'uv' is not installed or not on PATH."
echo ""
echo "Install uv by running:"
echo " curl -LsSf https://astral.sh/uv/install.sh | sh"
echo ""
echo "Then reopen your terminal and try again."
read -p "Press enter to exit..."
exit 1
fi
echo "Starting Corridor Key locally..."
echo "Target: $TARGET_PATH"
# Run via uv entry point (handles the virtual environment automatically)
uv run corridorkey wizard "$TARGET_PATH"
read -p "Press enter to close..."
================================================
FILE: Dockerfile
================================================
FROM ghcr.io/astral-sh/uv:0.7-python3.11-bookworm-slim
# Create non-root user upfront.
RUN useradd --create-home --uid 1000 appuser
RUN mkdir /app && chown appuser:appuser /app
WORKDIR /app
# Runtime dependencies for OpenCV/video I/O.
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
apt-get update && apt-get install -y --no-install-recommends \
ffmpeg \
git \
libgl1 \
libglib2.0-0 && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
USER appuser
# Install Python dependencies first for better layer caching.
COPY --chown=appuser:appuser pyproject.toml uv.lock ./
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --frozen --no-dev --no-install-project
# Copy project source.
COPY --chown=appuser:appuser . .
# Install the project itself (cheap, just sets up editable/entry points).
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --frozen --no-dev
# Enable OpenEXR support in OpenCV.
ENV OPENCV_IO_ENABLE_OPENEXR=1
ENTRYPOINT ["/app/.venv/bin/python", "corridorkey_cli.py"]
CMD ["--action", "list"]
================================================
FILE: IgnoredClips/.gitkeep
================================================
================================================
FILE: Install_CorridorKey_Linux_Mac.sh
================================================
#!/usr/bin/env bash
cd "$(dirname "$0")"
echo "==================================================="
echo " CorridorKey - MacOS/Linux Auto-Installer"
echo "==================================================="
echo ""
# Detect the Operating System
OS="$(uname -s)"
if [ "$OS" != "Darwin" ] && [ "$OS" != "Linux" ]; then
echo "[ERROR] Unsupported operating system: $OS"
read -p "Press [Enter] to exit..."
exit 1
fi
# 1. Check for uv — install it automatically if missing
if ! command -v uv >/dev/null 2>&1; then
echo "[INFO] uv is not installed. Installing now..."
curl -LsSf https://astral.sh/uv/install.sh | sh
if [ $? -ne 0 ]; then
echo "[ERROR] Failed to install uv. Please visit https://docs.astral.sh/uv/ for manual instructions."
read -p "Press [Enter] to exit..."
exit 1
fi
# uv installer adds to PATH, but the current terminal session
# doesn't see it yet. Add the default install location so we can continue.
export PATH="$HOME/.local/bin:$PATH"
if ! command -v uv >/dev/null 2>&1; then
echo "[ERROR] uv was installed but cannot be found on PATH."
echo "Please close this window, open a new terminal, and run this script again."
read -p "Press [Enter] to exit..."
exit 1
fi
echo "[INFO] uv installed successfully."
echo ""
fi
# 2. Install all dependencies
echo "[1/2] Installing Dependencies (This might take a while on first run)..."
echo " uv will automatically download Python if needed."
if [ "$OS" = "Darwin" ]; then
echo "[INFO] macOS detected. Installing with MLX support..."
uv sync --extra mlx
elif [ "$OS" = "Linux" ]; then
echo "[INFO] Linux detected. Installing with CUDA support..."
uv sync --extra cuda
fi
if [ $? -ne 0 ]; then
echo "[ERROR] uv sync failed. Please check the output above for details."
read -p "Press [Enter] to exit..."
exit 1
fi
# 3. Download Weights
echo ""
echo "[2/2] Downloading CorridorKey Model Weights..."
# Use -p to create the folder only if it doesn't exist
mkdir -p "CorridorKeyModule/checkpoints"
if [ ! -f "CorridorKeyModule/checkpoints/CorridorKey.pth" ]; then
echo "Downloading CorridorKey.pth..."
curl -L -o "CorridorKeyModule/checkpoints/CorridorKey.pth" "https://huggingface.co/nikopueringer/CorridorKey_v1.0/resolve/main/CorridorKey_v1.0.pth"
else
echo "CorridorKey.pth already exists!"
fi
echo ""
echo "==================================================="
echo " Setup Complete! You are ready to key!"
echo "==================================================="
read -p "Press [Enter] to close..."
================================================
FILE: Install_CorridorKey_Windows.bat
================================================
@echo off
TITLE CorridorKey Setup Wizard
echo ===================================================
echo CorridorKey - Windows Auto-Installer
echo ===================================================
echo.
:: 1. Check for uv — install it automatically if missing
where uv >nul 2>&1
if %errorlevel% equ 0 goto :uv_ready
echo [INFO] uv is not installed. Installing now...
powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"
if %errorlevel% neq 0 (
echo [ERROR] Failed to install uv. Please visit https://docs.astral.sh/uv/ for manual instructions.
pause
exit /b
)
:: uv installer adds to PATH via registry, but the current cmd session
:: doesn't see it yet. Add the default install location so we can continue.
set "PATH=%USERPROFILE%\.local\bin;%PATH%"
where uv >nul 2>&1
if %errorlevel% neq 0 (
echo [ERROR] uv was installed but cannot be found on PATH.
echo Please close this window, open a new terminal, and run this script again.
pause
exit /b
)
echo [INFO] uv installed successfully.
echo.
:uv_ready
:: 2. Install all dependencies (Python, venv, and packages are handled automatically by uv)
echo [1/2] Installing Dependencies (This might take a while on first run)...
echo uv will automatically download Python if needed.
uv sync --extra cuda
if %errorlevel% neq 0 (
echo [ERROR] uv sync failed. Please check the output above for details.
pause
exit /b
)
:: 3. Download Weights
echo.
echo [2/2] Downloading CorridorKey Model Weights...
if not exist "CorridorKeyModule\checkpoints" mkdir "CorridorKeyModule\checkpoints"
if not exist "CorridorKeyModule\checkpoints\CorridorKey.pth" (
echo Downloading CorridorKey.pth...
curl.exe -L -o "CorridorKeyModule\checkpoints\CorridorKey.pth" "https://huggingface.co/nikopueringer/CorridorKey_v1.0/resolve/main/CorridorKey_v1.0.pth"
) else (
echo CorridorKey.pth already exists!
)
echo.
echo ===================================================
echo Setup Complete! You are ready to key!
echo Drag and drop folders onto CorridorKey_DRAG_CLIPS_HERE_local.bat
echo ===================================================
pause
================================================
FILE: Install_GVM_Linux_Mac.sh
================================================
#!/usr/bin/env bash
cd "$(dirname "$0")"
# Set the Terminal window title
echo -n -e "\033]0;GVM Setup Wizard\007"
echo "==================================================="
echo " GVM (AlphaHint Generator) - Auto-Installer"
echo "==================================================="
echo ""
# Check that uv sync has been run (the .venv directory should exist)
# Note: I changed the name in the error message to match your Mac installer!
if [ ! -d ".venv" ]; then
echo "[ERROR] Project environment not found."
echo "Please run Install_CorridorKey_Linux_Mac.sh first!"
read -p "Press [Enter] to exit..."
exit 1
fi
# 1. Download Weights
echo "[1/1] Downloading GVM Model Weights (WARNING: Massive 80GB+ Download)..."
mkdir -p "gvm_core/weights"
echo "Downloading GVM weights from HuggingFace..."
uv run hf download geyongtao/gvm --local-dir "gvm_core/weights"
echo ""
echo "==================================================="
echo " GVM Setup Complete!"
echo "==================================================="
read -p "Press [Enter] to close..."
================================================
FILE: Install_GVM_Windows.bat
================================================
@echo off
TITLE GVM Setup Wizard
echo ===================================================
echo GVM (AlphaHint Generator) - Auto-Installer
echo ===================================================
echo.
:: Check that uv sync has been run (the .venv directory should exist)
if not exist ".venv" (
echo [ERROR] Project environment not found.
echo Please run Install_CorridorKey_Windows.bat first!
pause
exit /b
)
:: 1. Download Weights (all Python deps are already installed by uv sync)
echo [1/1] Downloading GVM Model Weights (WARNING: Massive 80GB+ Download)...
if not exist "gvm_core\weights" mkdir "gvm_core\weights"
echo Downloading GVM weights from HuggingFace...
uv run hf download geyongtao/gvm --local-dir gvm_core\weights
echo.
echo ===================================================
echo GVM Setup Complete!
echo ===================================================
pause
================================================
FILE: Install_VideoMaMa_Linux_Mac.sh
================================================
#!/usr/bin/env bash
cd "$(dirname "$0")"
# Set the Terminal window title
echo -n -e "\033]0;VideoMaMa Setup Wizard\007"
echo "==================================================="
echo " VideoMaMa (AlphaHint Generator) - Auto-Installer"
echo "==================================================="
echo ""
# Check that uv sync has been run (the .venv directory should exist)
# Note: I changed the name in the error message to match your Mac installer!
if [ ! -d ".venv" ]; then
echo "[ERROR] Project environment not found."
echo "Please run Install_CorridorKey_Linux_Mac.sh first!"
read -p "Press [Enter] to exit..."
exit 1
fi
# 1. Download Weights
echo "[1/1] Downloading VideoMaMa Model Weights..."
mkdir -p "VideoMaMaInferenceModule/checkpoints"
echo "Downloading VideoMaMa weights from HuggingFace..."
uv run hf download SammyLim/VideoMaMa --local-dir "VideoMaMaInferenceModule/checkpoints"
echo ""
echo "==================================================="
echo " VideoMaMa Setup Complete!"
echo "==================================================="
read -p "Press [Enter] to close..."
================================================
FILE: Install_VideoMaMa_Windows.bat
================================================
@echo off
TITLE VideoMaMa Setup Wizard
echo ===================================================
echo VideoMaMa (AlphaHint Generator) - Auto-Installer
echo ===================================================
echo.
:: Check that uv sync has been run (the .venv directory should exist)
if not exist ".venv" (
echo [ERROR] Project environment not found.
echo Please run Install_CorridorKey_Windows.bat first!
pause
exit /b
)
:: 1. Download Weights (all Python deps are already installed by uv sync)
echo [1/1] Downloading VideoMaMa Model Weights...
if not exist "VideoMaMaInferenceModule\checkpoints" mkdir "VideoMaMaInferenceModule\checkpoints"
echo Downloading VideoMaMa weights from HuggingFace...
uv run hf download SammyLim/VideoMaMa --local-dir VideoMaMaInferenceModule\checkpoints
echo.
echo ===================================================
echo VideoMaMa Setup Complete!
echo ===================================================
pause
================================================
FILE: LICENSE
================================================
CORRIDOR KEY LICENCE
=======================================================================
Version 1.0
Copyright (c) Corridor Digital. All rights reserved.
ADDITIONAL TERMS AND CONDITIONS
=======================================================================
This work is licensed under the Creative Commons
Attribution-NonCommercial-ShareAlike 4.0 International Public License
(CC BY-NC-SA 4.0), the full text of which is included below, subject
to the following additional terms and conditions. These additional
terms supplement the CC BY-NC-SA 4.0 licence and take precedence
where they conflict.
By exercising any rights to the Licensed Material, You accept and
agree to be bound by both these Additional Terms and the
CC BY-NC-SA 4.0 Public License.
1. PERMITTED USE
You may use this tool for any lawful purpose, including for
processing images as part of a commercial project, provided that
such use complies with the restrictions set out below and the
terms of the CC BY-NC-SA 4.0 licence.
2. RESTRICTIONS
In addition to the restrictions set out in the CC BY-NC-SA 4.0
licence, the following restrictions apply:
a. NO REPACKAGING OR RESALE
You may not repackage, redistribute, sublicense, or sell
this tool, in whole or in part, as a standalone product or
as part of a competing product.
b. NO PAID API OR INFERENCE SERVICES
You may not use this tool to provide inference as a paid
API service, whether directly or indirectly. This includes,
but is not limited to, offering access to this tool behind
a paywall, subscription model, or usage-based billing
system.
c. COMMERCIAL SOFTWARE INTEGRATION
If you operate a commercial software package or inference
service and wish to incorporate this tool into your
software, you must obtain a separate written agreement
from the Licensor. Please contact:
contact@corridordigital.com
3. ATTRIBUTION
a. In addition to the attribution requirements set out in
Section 3(a) of the CC BY-NC-SA 4.0 licence, You must
include the name "CorridorKey" in any attribution notice.
b. Any fork, derivative work, variation, improvement, or
subsequent release of this tool must retain the
"CorridorKey" name in a reasonably prominent position.
c. You may interchange "Corridor Key" for "CorridorKey" as
needed.
4. SHARE-ALIKE
Any variations, improvements, derivative works, or modified
versions of this tool that are publicly released must be
distributed under this same licence, including these Additional
Terms, or a licence with substantially equivalent terms.
5. DISCLAIMER
This licence is provided on an "as-is" basis. The Licensor makes
no representations or warranties regarding the enforceability or
legal sufficiency of these terms. You are encouraged to seek
independent legal advice if you have questions about the scope
or applicability of this licence.
6. CONTRIBUTIONS
By submitting any contribution (including source code,
documentation, or images) to this repository, You agree that:
a. Your contribution is provided under the terms of this
Corridor Key Licence.
b. You grant the Licensor (Corridor Digital) a perpetual,
worldwide, non-exclusive, no-charge, royalty-free,
irrevocable license to use, reproduce, prepare derivative
works of, publicly display, and sublicense your
contribution, including for commercial purposes.
c. You represent that you are the legal owner of the
contribution or have the authority to submit it under
these terms.
=======================================================================
https://creativecommons.org/licenses/by-nc-sa/4.0/
https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode.txt
=======================================================================
Attribution-NonCommercial-ShareAlike 4.0 International
=======================================================================
Creative Commons Corporation ("Creative Commons") is not a law firm and
does not provide legal services or legal advice. Distribution of
Creative Commons public licenses does not create a lawyer-client or
other relationship. Creative Commons makes its licenses and related
information available on an "as-is" basis. Creative Commons gives no
warranties regarding its licenses, any material licensed under their
terms and conditions, or any related information. Creative Commons
disclaims all liability for damages resulting from their use to the
fullest extent possible.
Using Creative Commons Public Licenses
Creative Commons public licenses provide a standard set of terms and
conditions that creators and other rights holders may use to share
original works of authorship and other material subject to copyright
and certain other rights specified in the public license below. The
following considerations are for informational purposes only, are not
exhaustive, and do not form part of our licenses.
Considerations for licensors: Our public licenses are
intended for use by those authorized to give the public
permission to use material in ways otherwise restricted by
copyright and certain other rights. Our licenses are
irrevocable. Licensors should read and understand the terms
and conditions of the license they choose before applying it.
Licensors should also secure all rights necessary before
applying our licenses so that the public can reuse the
material as expected. Licensors should clearly mark any
material not subject to the license. This includes other CC-
licensed material, or material used under an exception or
limitation to copyright. More considerations for licensors:
wiki.creativecommons.org/Considerations_for_licensors
Considerations for the public: By using one of our public
licenses, a licensor grants the public permission to use the
licensed material under specified terms and conditions. If
the licensor's permission is not necessary for any reason--for
example, because of any applicable exception or limitation to
copyright--then that use is not regulated by the license. Our
licenses grant only permissions under copyright and certain
other rights that a licensor has authority to grant. Use of
the licensed material may still be restricted for other
reasons, including because others have copyright or other
rights in the material. A licensor may make special requests,
such as asking that all changes be marked or described.
Although not required by our licenses, you are encouraged to
respect those requests where reasonable. More considerations
for the public:
wiki.creativecommons.org/Considerations_for_licensees
=======================================================================
Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
Public License
By exercising the Licensed Rights (defined below), You accept and agree
to be bound by the terms and conditions of this Creative Commons
Attribution-NonCommercial-ShareAlike 4.0 International Public License
("Public License"). To the extent this Public License may be
interpreted as a contract, You are granted the Licensed Rights in
consideration of Your acceptance of these terms and conditions, and the
Licensor grants You such rights in consideration of benefits the
Licensor receives from making the Licensed Material available under
these terms and conditions.
Section 1 -- Definitions.
a. Adapted Material means material subject to Copyright and Similar
Rights that is derived from or based upon the Licensed Material
and in which the Licensed Material is translated, altered,
arranged, transformed, or otherwise modified in a manner requiring
permission under the Copyright and Similar Rights held by the
Licensor. For purposes of this Public License, where the Licensed
Material is a musical work, performance, or sound recording,
Adapted Material is always produced where the Licensed Material is
synched in timed relation with a moving image.
b. Adapter's License means the license You apply to Your Copyright
and Similar Rights in Your contributions to Adapted Material in
accordance with the terms and conditions of this Public License.
c. BY-NC-SA Compatible License means a license listed at
creativecommons.org/compatiblelicenses, approved by Creative
Commons as essentially the equivalent of this Public License.
d. Copyright and Similar Rights means copyright and/or similar rights
closely related to copyright including, without limitation,
performance, broadcast, sound recording, and Sui Generis Database
Rights, without regard to how the rights are labeled or
categorized. For purposes of this Public License, the rights
specified in Section 2(b)(1)-(2) are not Copyright and Similar
Rights.
e. Effective Technological Measures means those measures that, in the
absence of proper authority, may not be circumvented under laws
fulfilling obligations under Article 11 of the WIPO Copyright
Treaty adopted on December 20, 1996, and/or similar international
agreements.
f. Exceptions and Limitations means fair use, fair dealing, and/or
any other exception or limitation to Copyright and Similar Rights
that applies to Your use of the Licensed Material.
g. License Elements means the license attributes listed in the name
of a Creative Commons Public License. The License Elements of this
Public License are Attribution, NonCommercial, and ShareAlike.
h. Licensed Material means the artistic or literary work, database,
or other material to which the Licensor applied this Public
License.
i. Licensed Rights means the rights granted to You subject to the
terms and conditions of this Public License, which are limited to
all Copyright and Similar Rights that apply to Your use of the
Licensed Material and that the Licensor has authority to license.
j. Licensor means the individual(s) or entity(ies) granting rights
under this Public License.
k. NonCommercial means not primarily intended for or directed towards
commercial advantage or monetary compensation. For purposes of
this Public License, the exchange of the Licensed Material for
other material subject to Copyright and Similar Rights by digital
file-sharing or similar means is NonCommercial provided there is
no payment of monetary compensation in connection with the
exchange.
l. Share means to provide material to the public by any means or
process that requires permission under the Licensed Rights, such
as reproduction, public display, public performance, distribution,
dissemination, communication, or importation, and to make material
available to the public including in ways that members of the
public may access the material from a place and at a time
individually chosen by them.
m. Sui Generis Database Rights means rights other than copyright
resulting from Directive 96/9/EC of the European Parliament and of
the Council of 11 March 1996 on the legal protection of databases,
as amended and/or succeeded, as well as other essentially
equivalent rights anywhere in the world.
n. You means the individual or entity exercising the Licensed Rights
under this Public License. Your has a corresponding meaning.
Section 2 -- Scope.
a. License grant.
1. Subject to the terms and conditions of this Public License,
the Licensor hereby grants You a worldwide, royalty-free,
non-sublicensable, non-exclusive, irrevocable license to
exercise the Licensed Rights in the Licensed Material to:
a. reproduce and Share the Licensed Material, in whole or
in part, for NonCommercial purposes only; and
b. produce, reproduce, and Share Adapted Material for
NonCommercial purposes only.
2. Exceptions and Limitations. For the avoidance of doubt, where
Exceptions and Limitations apply to Your use, this Public
License does not apply, and You do not need to comply with
its terms and conditions.
3. Term. The term of this Public License is specified in Section
6(a).
4. Media and formats; technical modifications allowed. The
Licensor authorizes You to exercise the Licensed Rights in
all media and formats whether now known or hereafter created,
and to make technical modifications necessary to do so. The
Licensor waives and/or agrees not to assert any right or
authority to forbid You from making technical modifications
necessary to exercise the Licensed Rights, including
technical modifications necessary to circumvent Effective
Technological Measures. For purposes of this Public License,
simply making modifications authorized by this Section 2(a)
(4) never produces Adapted Material.
5. Downstream recipients.
a. Offer from the Licensor -- Licensed Material. Every
recipient of the Licensed Material automatically
receives an offer from the Licensor to exercise the
Licensed Rights under the terms and conditions of this
Public License.
b. Additional offer from the Licensor -- Adapted Material.
Every recipient of Adapted Material from You
automatically receives an offer from the Licensor to
exercise the Licensed Rights in the Adapted Material
under the conditions of the Adapter's License You apply.
c. No downstream restrictions. You may not offer or impose
any additional or different terms or conditions on, or
apply any Effective Technological Measures to, the
Licensed Material if doing so restricts exercise of the
Licensed Rights by any recipient of the Licensed
Material.
6. No endorsement. Nothing in this Public License constitutes or
may be construed as permission to assert or imply that You
are, or that Your use of the Licensed Material is, connected
with, or sponsored, endorsed, or granted official status by,
the Licensor or others designated to receive attribution as
provided in Section 3(a)(1)(A)(i).
b. Other rights.
1. Moral rights, such as the right of integrity, are not
licensed under this Public License, nor are publicity,
privacy, and/or other similar personality rights; however, to
the extent possible, the Licensor waives and/or agrees not to
assert any such rights held by the Licensor to the limited
extent necessary to allow You to exercise the Licensed
Rights, but not otherwise.
2. Patent and trademark rights are not licensed under this
Public License.
3. To the extent possible, the Licensor waives any right to
collect royalties from You for the exercise of the Licensed
Rights, whether directly or through a collecting society
under any voluntary or waivable statutory or compulsory
licensing scheme. In all other cases the Licensor expressly
reserves any right to collect such royalties, including when
the Licensed Material is used other than for NonCommercial
purposes.
Section 3 -- License Conditions.
Your exercise of the Licensed Rights is expressly made subject to the
following conditions.
a. Attribution.
1. If You Share the Licensed Material (including in modified
form), You must:
a. retain the following if it is supplied by the Licensor
with the Licensed Material:
i. identification of the creator(s) of the Licensed
Material and any others designated to receive
attribution, in any reasonable manner requested by
the Licensor (including by pseudonym if
designated);
ii. a copyright notice;
iii. a notice that refers to this Public License;
iv. a notice that refers to the disclaimer of
warranties;
v. a URI or hyperlink to the Licensed Material to the
extent reasonably practicable;
b. indicate if You modified the Licensed Material and
retain an indication of any previous modifications; and
c. indicate the Licensed Material is licensed under this
Public License, and include the text of, or the URI or
hyperlink to, this Public License.
2. You may satisfy the conditions in Section 3(a)(1) in any
reasonable manner based on the medium, means, and context in
which You Share the Licensed Material. For example, it may be
reasonable to satisfy the conditions by providing a URI or
hyperlink to a resource that includes the required
information.
3. If requested by the Licensor, You must remove any of the
information required by Section 3(a)(1)(A) to the extent
reasonably practicable.
b. ShareAlike.
In addition to the conditions in Section 3(a), if You Share
Adapted Material You produce, the following conditions also apply.
1. The Adapter's License You apply must be a Creative Commons
license with the same License Elements, this version or
later, or a BY-NC-SA Compatible License.
2. You must include the text of, or the URI or hyperlink to, the
Adapter's License You apply. You may satisfy this condition
in any reasonable manner based on the medium, means, and
context in which You Share Adapted Material.
3. You may not offer or impose any additional or different terms
or conditions on, or apply any Effective Technological
Measures to, Adapted Material that restrict exercise of the
rights granted under the Adapter's License You apply.
Section 4 -- Sui Generis Database Rights.
Where the Licensed Rights include Sui Generis Database Rights that
apply to Your use of the Licensed Material:
a. for the avoidance of doubt, Section 2(a)(1) grants You the right
to extract, reuse, reproduce, and Share all or a substantial
portion of the contents of the database for NonCommercial purposes
only;
b. if You include all or a substantial portion of the database
contents in a database in which You have Sui Generis Database
Rights, then the database in which You have Sui Generis Database
Rights (but not its individual contents) is Adapted Material,
including for purposes of Section 3(b); and
c. You must comply with the conditions in Section 3(a) if You Share
all or a substantial portion of the contents of the database.
For the avoidance of doubt, this Section 4 supplements and does not
replace Your obligations under this Public License where the Licensed
Rights include other Copyright and Similar Rights.
Section 5 -- Disclaimer of Warranties and Limitation of Liability.
a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
c. The disclaimer of warranties and limitation of liability provided
above shall be interpreted in a manner that, to the extent
possible, most closely approximates an absolute disclaimer and
waiver of all liability.
Section 6 -- Term and Termination.
a. This Public License applies for the term of the Copyright and
Similar Rights licensed here. However, if You fail to comply with
this Public License, then Your rights under this Public License
terminate automatically.
b. Where Your right to use the Licensed Material has terminated under
Section 6(a), it reinstates:
1. automatically as of the date the violation is cured, provided
it is cured within 30 days of Your discovery of the
violation; or
2. upon express reinstatement by the Licensor.
For the avoidance of doubt, this Section 6(b) does not affect any
right the Licensor may have to seek remedies for Your violations
of this Public License.
c. For the avoidance of doubt, the Licensor may also offer the
Licensed Material under separate terms or conditions or stop
distributing the Licensed Material at any time; however, doing so
will not terminate this Public License.
d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
License.
Section 7 -- Other Terms and Conditions.
a. The Licensor shall not be bound by any additional or different
terms or conditions communicated by You unless expressly agreed.
b. Any arrangements, understandings, or agreements regarding the
Licensed Material not stated herein are separate from and
independent of the terms and conditions of this Public License.
Section 8 -- Interpretation.
a. For the avoidance of doubt, this Public License does not, and
shall not be interpreted to, reduce, limit, restrict, or impose
conditions on any use of the Licensed Material that could lawfully
be made without permission under this Public License.
b. To the extent possible, if any provision of this Public License is
deemed unenforceable, it shall be automatically reformed to the
minimum extent necessary to make it enforceable. If the provision
cannot be reformed, it shall be severed from this Public License
without affecting the enforceability of the remaining terms and
conditions.
c. No term or condition of this Public License will be waived and no
failure to comply consented to unless expressly agreed to by the
Licensor.
d. Nothing in this Public License constitutes or may be interpreted
as a limitation upon, or waiver of, any privileges and immunities
that apply to the Licensor or You, including from the legal
processes of any jurisdiction or authority.
=======================================================================
Creative Commons is not a party to its public
licenses. Notwithstanding, Creative Commons may elect to apply one of
its public licenses to material it publishes and in those instances
will be considered the “Licensor.” The text of the Creative Commons
public licenses is dedicated to the public domain under the CC0 Public
Domain Dedication. Except for the limited purpose of indicating that
material is shared under a Creative Commons public license or as
otherwise permitted by the Creative Commons policies published at
creativecommons.org/policies, Creative Commons does not authorize the
use of the trademark "Creative Commons" or any other trademark or logo
of Creative Commons without its prior written consent including,
without limitation, in connection with any unauthorized modifications
to any of its public licenses or any other arrangements,
understandings, or agreements concerning use of licensed material. For
the avoidance of doubt, this paragraph does not form part of the
public licenses.
Creative Commons may be contacted at creativecommons.org.
================================================
FILE: Output/.gitkeep
================================================
================================================
FILE: README.md
================================================
# CorridorKey
https://github.com/user-attachments/assets/1fb27ea8-bc91-4ebc-818f-5a3b5585af08
When you film something against a green screen, the edges of your subject inevitably blend with the green background. This creates pixels that are a mix of your subject's color and the green screen's color. Traditional keyers struggle to untangle these colors, forcing you to spend hours building complex edge mattes or manually rotoscoping. Even modern "AI Roto" solutions typically output a harsh binary mask, completely destroying the delicate, semi-transparent pixels needed for a realistic composite.
I built CorridorKey to solve this *unmixing* problem.
You input a raw green screen frame, and the neural network completely separates the foreground object from the green screen. For every single pixel, even the highly transparent ones like motion blur or out-of-focus edges, the model predicts the true, un-multiplied straight color of the foreground element, alongside a clean, linear alpha channel. It doesn't just guess what is opaque and what is transparent; it actively reconstructs the color of the foreground object as if the green screen was never there.
No more fighting with garbage mattes or agonizing over "core" vs "edge" keys. Give CorridorKey a hint of what you want, and it separates the light for you.
## Alert!
This is a brand new release, I'm sure you will discover many ways it can be improved! I invite everyone to help. Join us on the "Corridor Creates" Discord to share ideas, work, forks, etc! https://discord.gg/zvwUrdWXJm
If you want an easy-install, artist-friendly user interface version of CorridorKey, check out [EZ-CorridorKey](https://github.com/edenaion/EZ-CorridorKey)
This project uses [uv](https://docs.astral.sh/uv/) to manage dependencies — it handles Python installation, virtual environments, and packages all in one step, so you don't need to worry about any of that. just run the appropriate install script for your OS.
Naturally, I have not tested everything. If you encounter errors, please consider patching the code as needed and submitting a pull request.
## Features
* **Physically Accurate Unmixing:** Clean extraction of straight color foreground and linear alpha channels, preserving hair, motion blur, and translucency.
* **Resolution Independent:** The engine dynamically scales inference to handle 4K plates while predicting using its native 2048x2048 high-fidelity backbone.
* **VFX Standard Outputs:** Natively reads and writes 16-bit and 32-bit Linear float EXR files, preserving true color math for integration in Nuke, Fusion, or Resolve.
* **Auto-Cleanup:** Includes a morphological cleanup system to automatically prune any tracking markers or tiny background features that slip through CorridorKey's detection.
## Hardware Requirements
This project was designed and built on a Linux workstation (Puget Systems PC) equipped with an NVIDIA RTX Pro 6000 with 96GB of VRAM. The community is ACTIVELY optimizing it for consumer GPUS.
The most recent build should work on computers with 6-8 gig of VRAM, and it can run on most M1+ Mac systems with unified memory. Yes, it might even work on your old Macbook pro. Let us know on the Discord!
* **Windows Users:** To run GPU acceleration natively on Windows, your system MUST have NVIDIA drivers that support **CUDA 12.8 or higher** installed. If your drivers only support older CUDA versions, the installer will likely fallback to the CPU.
* **GVM (Optional):** Requires approximately **80 GB of VRAM** and utilizes massive Stable Video Diffusion models.
* **VideoMaMa (Optional):** Natively requires a massive chunk of VRAM as well (originally 80GB+). While the community has tweaked the architecture to run at less than 24GB, those extreme memory optimizations have not yet been fully implemented in this repository.
* **BiRefNet (Optional):** Lightweight AlphaHint generator option.
Because GVM and VideoMaMa have huge model file sizes and extreme hardware requirements, installing their modules is completely optional. You can always provide your own Alpha Hints generated from your editing program, BiRefNet, or any other method. The better the AlphaHint, the better the result.
## Getting Started
### 1. Installation
This project uses **[uv](https://docs.astral.sh/uv/)** to manage Python and all dependencies. uv is a fast, modern replacement for pip that automatically handles Python versions, virtual environments, and package installation in a single step. You do **not** need to install Python yourself — uv does it for you.
**For Windows Users (Automated):**
1. Clone or download this repository to your local machine.
2. Double-click `Install_CorridorKey_Windows.bat`. This will automatically install uv (if needed), set up your Python environment, install all dependencies, and download the CorridorKey model.
> **Note:** If this is the first time installing uv, any terminal windows you already had open won't see it. The installer script handles the current window automatically, but if you open a new terminal and get "'uv' is not recognized", just close and reopen that terminal.
3. (Optional) Double-click `Install_GVM_Windows.bat` and `Install_VideoMaMa_Windows.bat` to download the heavy optional Alpha Hint generator weights.
**For Linux / Mac Users (Automated):**
1. Clone or download this repository to your local machine.
2. Open terminal and write `bash`. Put a space after writing `bash`.
3. Drag and drop `Install_CorridorKey_Linux_Mac.sh` into the terminal. Then press enter.
4. (Optional) Do the 2. step again. But now drag and drop `Install_GVM_Linux_Mac.sh` and `Install_VideoMaMa_Linux_Mac.sh` to download the heavy optional Alpha Hint generator weights.
**For Linux / Mac Users (Manual):**
1. Clone or download this repository to your local machine.
2. Install uv if you don't have it:
```bash
curl -LsSf https://astral.sh/uv/install.sh | sh
```
3. Install all dependencies (uv will download Python 3.10+ automatically if needed):
```bash
uv sync # CPU/MPS (default — works everywhere)
uv sync --extra cuda # CUDA GPU acceleration (Linux/Windows)
uv sync --extra mlx # Apple Silicon MLX acceleration
```
4. **Download the Models:**
* **CorridorKey v1.0 Model (~300MB):** Downloads automatically on first run. If no `.pth` file is found in `CorridorKeyModule/checkpoints/`, the engine fetches it from [CorridorKey's HuggingFace](https://huggingface.co/nikopueringer/CorridorKey_v1.0) and saves it as `CorridorKey.pth`. No manual download needed.
* **GVM Weights (Optional):** [HuggingFace: geyongtao/gvm](https://huggingface.co/geyongtao/gvm)
* Download using the CLI: `uv run hf download geyongtao/gvm --local-dir gvm_core/weights`
* **VideoMaMa Weights (Optional):** [HuggingFace: SammyLim/VideoMaMa](https://huggingface.co/SammyLim/VideoMaMa)
* Download the VideoMaMa fine-tuned weights:
```
uv run hf download SammyLim/VideoMaMa --local-dir VideoMaMaInferenceModule/checkpoints/VideoMaMa
```
* VideoMaMa also requires the Stable Video Diffusion base model (VAE + image encoder only, ~2.5GB). Accept the license at [stabilityai/stable-video-diffusion-img2vid-xt](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt), then:
```
uv run hf download stabilityai/stable-video-diffusion-img2vid-xt \
--local-dir VideoMaMaInferenceModule/checkpoints/stable-video-diffusion-img2vid-xt \
--include "feature_extractor/*" "image_encoder/*" "vae/*" "model_index.json"
```
* VideoMaMa is an amazing project, please go star their [repo](https://github.com/cvlab-kaist/VideoMaMa) and show them some support!
### 2. How it Works
CorridorKey requires two inputs to process a frame:
1. **The Original RGB Image:** The to-be-processed green screen footage. This requires the sRGB color gamut (interchangeable with REC709 gamut), and the engine can ingest either an sRGB gamma or Linear gamma curve.
2. **A Coarse Alpha Hint:** A rough black-and-white mask that generally isolates the subject. This does *not* need to be precise. It can be generated by you with a rough chroma key or AI roto.
I've had the best results using GVM or VideoMaMa to create the AlphaHint, so I've repackaged those projects and integrated them here as optional modules inside `clip_manager.py`. Here is how they compare:
* **GVM:** Completely automatic and requires no additional input. It works exceptionally well for people, but can struggle with inanimate objects.
* **VideoMaMa:** Requires you to provide a rough VideoMamaMaskHint (often drawn by hand or AI) telling it what you want to key. If you choose to use this, place your mask hint in the `VideoMamaMaskHint/` folder that the wizard creates for your shot. VideoMaMa results are spectacular and can be controlled more easily than GVM due to this mask hint.
* **Please** go show the creators of these projects some love and star their repos. [VideoMaMa](https://github.com/cvlab-kaist/VideoMaMa) and [GVM](https://github.com/aim-uofa/GVM)
Perhaps in the future, I will implement other generators for the AlphaHint! In the meantime, the better your Alpha Hint, the better CorridorKey's final result will be. Experiment with different amounts of mask erosion or feathering. The model was trained on coarse, blurry, eroded masks, and is exceptional at filling in details from the hint. However, it is generally less effective at subtracting unwanted mask details if your Alpha Hint is expanded too far.
Please give feedback and share your results!
### Docker (Linux + NVIDIA GPU)
If you prefer not to install dependencies locally, you can run CorridorKey in Docker.
Prerequisites:
- Docker Engine + Docker Compose plugin installed.
- NVIDIA driver installed on the host (Linux), with CUDA compatibility for the PyTorch CUDA 12.6 wheels used by this project.
- NVIDIA Container Toolkit installed and configured for Docker (`nvidia-smi` should work on host, and `docker run --rm --gpus all nvidia/cuda:12.6.3-runtime-ubuntu22.04 nvidia-smi` should succeed).
1. Build the image:
```bash
docker build -t corridorkey:latest .
```
2. Run an action directly (example: inference):
```bash
docker run --rm -it --gpus all \
-e OPENCV_IO_ENABLE_OPENEXR=1 \
-v "$(pwd)/ClipsForInference:/app/ClipsForInference" \
-v "$(pwd)/Output:/app/Output" \
-v "$(pwd)/CorridorKeyModule/checkpoints:/app/CorridorKeyModule/checkpoints" \
-v "$(pwd)/gvm_core/weights:/app/gvm_core/weights" \
-v "$(pwd)/VideoMaMaInferenceModule/checkpoints:/app/VideoMaMaInferenceModule/checkpoints" \
corridorkey:latest run_inference --device cuda
```
3. Docker Compose (recommended for repeat runs):
```bash
docker compose build
docker compose --profile gpu run --rm corridorkey run_inference --device cuda
docker compose --profile gpu run --rm corridorkey list
docker compose --profile cpu run --rm corridorkey-cpu run_inference --device cpu
```
4. Optional: pin to specific GPU(s) for multi-GPU workstations:
```bash
NVIDIA_VISIBLE_DEVICES=0 docker compose --profile gpu run --rm corridorkey list
NVIDIA_VISIBLE_DEVICES=1,2 docker compose --profile gpu run --rm corridorkey run_inference --device cuda
```
Notes:
- You still need to place model weights in the same folders used by native runs (mounted above).
- The container does not include kernel GPU drivers; those always come from the host. The image provides user-space dependencies and relies on Docker's NVIDIA runtime to pass through driver libraries/devices.
- The wizard works too, but use a path inside the container, for example:
```bash
docker run --rm -it --gpus all \
-e OPENCV_IO_ENABLE_OPENEXR=1 \
-v "$(pwd)/ClipsForInference:/app/ClipsForInference" \
-v "$(pwd)/Output:/app/Output" \
-v "$(pwd)/CorridorKeyModule/checkpoints:/app/CorridorKeyModule/checkpoints" \
-v "$(pwd)/gvm_core/weights:/app/gvm_core/weights" \
-v "$(pwd)/VideoMaMaInferenceModule/checkpoints:/app/VideoMaMaInferenceModule/checkpoints" \
corridorkey:latest wizard --win_path /app/ClipsForInference
docker compose --profile gpu run --rm corridorkey wizard --win_path /app/ClipsForInference
```
### 3. Usage: The Command Line Wizard
For the easiest experience, use the provided launcher scripts. These scripts launch a prompt-based configuration wizard in your terminal.
* **Windows:** Drag-and-drop a video file or folder onto `CorridorKey_DRAG_CLIPS_HERE_local.bat` (Note: Only launch via Drag-and-Drop or CMD. Double-clicking the `.bat` directly will throw an error).
* **Linux / Mac:** Run or drag-and-drop a video file or folder onto `./CorridorKey_DRAG_CLIPS_HERE_local.sh`.
* - Or write `bash` again in terminal. Put a space after and then drag-and-drop `CorridorKey_DRAG_CLIPS_HERE_local.sh` and your clip folder together into terminal, respectively. Then press enter.
**Workflow Steps:**
1. **Launch:** You can drag-and-drop a single loose video file (like an `.mp4`), a shot folder containing image sequences, or even a master "batch" folder containing multiple different shots all at once onto the launcher script.
2. **Organization:** The wizard will detect what you dragged in. If you dropped loose video files or unorganized folders, the first prompt will ask if you want it to organize your clips into the proper structure.
* If you say Yes, the script will automatically create a shot folder, move your footage into an `Input/` sub-folder, and generate empty `AlphaHint/` and `VideoMamaMaskHint/` folders for you. This structure is required for the engine to pair your hints and footage correctly!
3. **Generate Hints (Optional):** If the wizard detects your shots are missing an `AlphaHint`, it will ask if you want to generate them automatically using the repackaged GVM or VideoMaMa modules.
4. **Configure:** Once your clips have both Inputs and AlphaHints, select "Process Ready Clips". The wizard will prompt you to configure the run:
* **Gamma Space:** Tell the engine if your sequence uses a Linear or sRGB gamma curve.
* **Despill Strength:** This is a traditional despill filter (0-10), if you wish to have it baked into the output now as opposed to applying it in your comp later.
* **Auto-Despeckle:** Toggle automatic cleanup and define the size threshold. This isn't just for tracking dots, it removes any small, disconnected islands of pixels.
* **Refiner Strength:** Use the default (1.0) unless you are experimenting with extreme detail pushing.
5. **Result:** The engine will generate several folders inside your shot directory:
* `/Matte`: The raw Linear Alpha channel (EXR).
* `/FG`: The raw Straight Foreground Color Object. (Note: The engine natively computes this in the sRGB gamut. You must manually convert this pass to linear gamma before being combined with the alpha in your compositing program).
* `/Processed`: An RGBA image containing the Linear Foreground premultiplied against the Linear Alpha (EXR). This pass exists so you can immediately drop the footage into Premiere/Resolve for a quick preview without dealing with complex premultiplication routing. However, if you want more control over your image, working with the raw FG and Matte outputs will give you that.
* `/Comp`: A simple preview of the key composited over a checkerboard (PNG).
## But What About Training and Datasets?
If enough people find this project interesting I'll get the training program and datasets uploaded so we can all really go to town making the absolute best keyer fine tunes! Just hit me with some messages on the Corridor Creates discord or here. If enough people lock in, I'll get this stuff packaged up. Hardware requirements are beefy and the gigabytes are plentiful so I don't want to commit the time unless there's demand.
## Device Selection
By default, CorridorKey auto-detects the best available compute device: **CUDA > MPS > CPU**.
**Override via CLI flag:**
```bash
uv run python clip_manager.py --action wizard --win_path "V:\..." --device mps
uv run python clip_manager.py --action run_inference --device cpu
```
**Override via environment variable:**
```bash
export CORRIDORKEY_DEVICE=cpu
uv run python clip_manager.py --action wizard --win_path "V:\..."
```
Priority: `--device` flag > `CORRIDORKEY_DEVICE` env var > auto-detect.
### Apple Silicon / MPS Troubleshooting
**Confirm MPS is active:** Run with verbose logging to see which device was selected:
```bash
uv run python clip_manager.py --action list 2>&1 | grep -i "device\|backend\|mps"
```
**MPS operator errors** (`NotImplementedError: ... not implemented for 'MPS'`): Some PyTorch operations are not yet supported on MPS. Enable CPU fallback for those ops:
```bash
export PYTORCH_ENABLE_MPS_FALLBACK=1
uv run python corridorkey_cli.py wizard --win_path "/path/to/clips"
```
**Silent CPU fallback**: If MPS silently falls back to CPU without this variable, the run will be much slower. Setting `PYTORCH_ENABLE_MPS_FALLBACK=1` in your shell profile (`~/.zshrc`) ensures it is always active.
**Use native MLX instead of PyTorch MPS:** MLX avoids PyTorch's MPS layer entirely and typically runs faster on Apple Silicon. See the [Backend Selection](#backend-selection) section below for setup steps.
## Backend Selection
CorridorKey supports two inference backends:
- **Torch** (default on Linux/Windows) — CUDA, MPS, or CPU
- **MLX** (Apple Silicon) — native Metal acceleration, no Torch overhead
Resolution: `--backend` flag > `CORRIDORKEY_BACKEND` env var > auto-detect.
Auto mode prefers MLX on Apple Silicon when available.
**Override via CLI flag (corridorkey_cli.py):**
```bash
uv run python corridorkey_cli.py wizard --win_path "/path/to/clips" --backend mlx
uv run python corridorkey_cli.py run_inference --backend torch
```
### MLX Setup (Apple Silicon)
1. Install the MLX backend:
```bash
uv sync --extra mlx
```
2. Obtain the MLX weights (`.safetensors`) — pick **one** option:
**Option A — Download pre-converted weights (simplest):**
```bash
# Download weights from GitHub Releases into a local cache directory
uv run python -m corridorkey_mlx weights download
# Print the cached path, then copy to the checkpoints folder
WEIGHTS=$(uv run python -m corridorkey_mlx weights download --print-path)
cp "$WEIGHTS" CorridorKeyModule/checkpoints/corridorkey_mlx.safetensors
```
**Option B — Convert from an existing `.pth` checkpoint:**
```bash
# Clone the MLX repo (contains the conversion script)
git clone https://github.com/nikopueringer/corridorkey-mlx.git
cd corridorkey-mlx
uv sync
# Convert (point --checkpoint at your CorridorKey.pth)
uv run python scripts/convert_weights.py \
--checkpoint ../CorridorKeyModule/checkpoints/CorridorKey_v1.0.pth \
--output ../CorridorKeyModule/checkpoints/corridorkey_mlx.safetensors
cd ..
```
Either way the final file must be at:
```
CorridorKeyModule/checkpoints/corridorkey_mlx.safetensors
```
3. Run with auto-detection or explicit backend:
```bash
CORRIDORKEY_BACKEND=mlx uv run python clip_manager.py --action run_inference
```
MLX uses img_size=2048 by default (same as Torch).
### Troubleshooting
- **"No .safetensors checkpoint found"** — place MLX weights in `CorridorKeyModule/checkpoints/`
- **"corridorkey_mlx not installed"** — run `uv sync --extra mlx`
- **"MLX requires Apple Silicon"** — MLX only works on M1+ Macs
- **Auto picked Torch unexpectedly** — set `CORRIDORKEY_BACKEND=mlx` explicitly
## Advanced Usage
For developers looking for more details on the specifics of what is happening in the CorridorKey engine, check out the README in the `/CorridorKeyModule` folder. We also have a dedicated handover document outlining the pipeline architecture for AI assistants in `/docs/LLM_HANDOVER.md`.
You can also explore the full, auto-generated codebase documentation on [DeepWiki](https://deepwiki.com/nikopueringer/CorridorKey).
### Running Tests
The project includes unit tests for the color math and compositing pipeline. No GPU or model weights required — tests run in a few seconds on any machine.
```bash
uv sync --group dev # install test dependencies (pytest)
uv run pytest # run all tests
uv run pytest -v # verbose output (shows each test name)
```
## CorridorKey Licensing and Permissions
Use this tool for whatever you'd like, including for processing images as part of a commercial project! You MAY NOT repackage this tool and sell it, and any variations or improvements of this tool that are released must remain under the same license, and must include the name Corridor Key.
You MAY NOT offer inference with this model as a paid API service. If you run a commercial software package or inference service and wish to incoporate this tool into your software, shoot us an email to work out an agreement! I promise we're easy to work with. contact@corridordigital.com. Outside of the stipulations listed above, this license is effectively a variation of [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License (CC BY-NC-SA 4.0)](https://creativecommons.org/licenses/by-nc-sa/4.0/)
Please keep the Corridor Key name in any future forks or releases!
## Community Extensions
* [CorridorKeyOpenVINO](https://github.com/daniil-lyakhov/CorridorKeyOpenVINO) - Run the CorridorKey model quickly on Intel hardware with the OpenVINO inference framework.
## Acknowledgements and Licensing
CorridorKey integrates several open-source modules for Alpha Hint generation. We would like to explicitly credit and thank the following research teams:
* **Generative Video Matting (GVM):** Developed by the Advanced Intelligent Machines (AIM) research team at Zhejiang University. The GVM code and models are heavily utilized in the `gvm_core` module. Their work is licensed under the [2-clause BSD License (BSD-2-Clause)](https://opensource.org/license/bsd-2-clause). You can find their source repository here: [aim-uofa/GVM](https://github.com/aim-uofa/GVM). Give them a star!
* **VideoMaMa:** Developed by the CVLAB at KAIST. The VideoMaMa architecture is utilized within the `VideoMaMaInferenceModule`. Their code is released under the [Creative Commons Attribution-NonCommercial 4.0 International License (CC BY-NC 4.0)](https://creativecommons.org/licenses/by-nc/4.0/), and their specific foundation model checkpoints (`dino_projection_mlp.pth`, `unet/*`) are subject to the [Stability AI Community License](https://stability.ai/license). You can find their source repository here: [cvlab-kaist/VideoMaMa](https://github.com/cvlab-kaist/VideoMaMa). Give them a star!
By using these optional modules, you agree to abide by their respective Non-Commercial licenses. Please review their repositories for full terms.
================================================
FILE: RunGVMOnly.sh
================================================
#!/usr/bin/env bash
# Ensure script stops on error
set -e
# Path to script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Enable OpenEXR Support
export OPENCV_IO_ENABLE_OPENEXR=1
echo "Starting Coarse Alpha Generation..."
echo "Scanning ClipsForInference..."
# Run via uv entry point (handles the virtual environment automatically)
uv run corridorkey generate-alphas
echo "Done."
================================================
FILE: RunInferenceOnly.sh
================================================
#!/usr/bin/env bash
# Ensure script stops on error
set -e
# Path to script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Enable OpenEXR Support
export OPENCV_IO_ENABLE_OPENEXR=1
echo "Starting CorridorKey Inference..."
echo "Scanning ClipsForInference for Ready Clips (Input + Alpha)..."
# Run via uv entry point (handles the virtual environment automatically)
uv run corridorkey run-inference
echo "Inference Complete."
================================================
FILE: VideoMaMaInferenceModule/LICENSE.md
================================================
# VideoMaMa Licensing and Acknowledgements
This module (`VideoMaMaInferenceModule`) contains repackaged code and integrations from the **Video Masked Modeling (VideoMaMa)** project developed by the CVLAB at KAIST.
## Original Repository
* **GitHub:** [https://github.com/cvlab-kaist/VideoMaMa](https://github.com/cvlab-kaist/VideoMaMa)
* **HuggingFace:** [https://huggingface.co/SammyLim/VideoMaMa](https://huggingface.co/SammyLim/VideoMaMa)
## License
The VideoMaMa codebase is released under the **Creative Commons Attribution-NonCommercial 4.0 International License (CC BY-NC 4.0)**.
To view a copy of this license, visit: [http://creativecommons.org/licenses/by-nc/4.0/](http://creativecommons.org/licenses/by-nc/4.0/)
Additionally, the model checkpoints are subject to the **Stability AI Community License**.
To view a copy of this license, visit: [https://stability.ai/license](https://stability.ai/license)
By utilizing this module and downloading the associated weights, you are subject to these Non-Commercial and Community licenses.
================================================
FILE: VideoMaMaInferenceModule/README.md
================================================
# VideoMaMa Inference Module
This module provides a standalone interface for running VideoMaMa inference.
## Usage
```python
import sys
# Ensure the parent directory of this module is in sys.path
sys.path.append("/path/to/parent/directory")
from VideoMaMa_Inference_Module import load_videomama_model, run_inference, extract_frames_from_video, save_video
# 1. Load Model
# By default, it loads checkpoints from the local 'checkpoints/' directory inside the module.
# Ensure you have copied 'stable-video-diffusion-img2vid-xt' and 'VideoMaMa' into 'checkpoints/'.
pipeline = load_videomama_model(device="cuda")
# Alternatively, specify custom paths:
# pipeline = load_videomama_model(base_model_path="/path/to/base", unet_checkpoint_path="/path/to/unet", device="cuda")
# 2. Prepare Inputs
# You need a list of RGB frames and a list of mask frames (grayscale)
# Helper function to extract from video:
video_path = "input_video.mp4"
input_frames, fps = extract_frames_from_video(video_path, max_frames=24)
# Load your masks (e.g. from file or other process)
# masks = [ ... list of numpy arrays ... ]
# Ensure len(masks) == len(input_frames)
# 3. Run Inference
output_frames = run_inference(pipeline, input_frames, masks)
# 4. Save Output
save_video(output_frames, "output.mp4", fps)
```
## Requirements
Install dependencies listed in `requirements.txt`.
```bash
pip install -r requirements.txt
```
================================================
FILE: VideoMaMaInferenceModule/__init__.py
================================================
from .inference import load_videomama_model, run_inference, extract_frames_from_video, save_video
from .pipeline import VideoInferencePipeline
__all__ = [
"load_videomama_model",
"run_inference",
"extract_frames_from_video",
"save_video",
"VideoInferencePipeline"
]
================================================
FILE: VideoMaMaInferenceModule/checkpoints/.gitkeep
================================================
================================================
FILE: VideoMaMaInferenceModule/inference.py
================================================
"""
VideoMaMa Inference Module
Provides functions to load the model and run inference on video inputs.
"""
import os
import sys
import torch
import cv2
import numpy as np
from PIL import Image
from typing import List, Union, Optional
from pathlib import Path
# Add current directory to path so that pipeline.py's intra-package imports
# (e.g. "from pipeline import ...") resolve when this module is imported from
# outside the VideoMaMaInferenceModule directory. This is a workaround for the
# module's original structure — a cleaner fix would convert to proper relative
# imports throughout.
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
sys.path.append(current_dir)
from .pipeline import VideoInferencePipeline
def load_videomama_model(base_model_path: Optional[str] = None, unet_checkpoint_path: Optional[str] = None, device: str = "cpu") -> VideoInferencePipeline:
"""
Load VideoMaMa pipeline with pretrained weights.
Args:
base_model_path (str, optional): Path to the base Stable Video Diffusion model.
Defaults to 'checkpoints/stable-video-diffusion-img2vid-xt' in module dir.
unet_checkpoint_path (str, optional): Path to the fine-tuned UNet checkpoint.
Defaults to 'checkpoints/VideoMaMa' in module dir.
device (str): Device to run on ("cuda" or "cpu").
Returns:
VideoInferencePipeline: Loaded pipeline instance.
"""
# Default to local checkpoints if not provided
if base_model_path is None:
base_model_path = os.path.join(current_dir, "checkpoints", "stable-video-diffusion-img2vid-xt")
if unet_checkpoint_path is None:
unet_checkpoint_path = os.path.join(current_dir, "checkpoints", "VideoMaMa")
print(f"Loading Base model from {base_model_path}...")
print(f"Loading VideoMaMa UNet from {unet_checkpoint_path}...")
# Check if paths exist
if not os.path.exists(base_model_path):
raise FileNotFoundError(f"Base model path not found: {base_model_path}")
if not os.path.exists(unet_checkpoint_path):
raise FileNotFoundError(f"UNet checkpoint path not found: {unet_checkpoint_path}")
pipeline = VideoInferencePipeline(
base_model_path=base_model_path,
unet_checkpoint_path=unet_checkpoint_path,
weight_dtype=torch.float16, # Use float16 for inference by default
device=device
)
print("VideoMaMa pipeline loaded successfully!")
return pipeline
def extract_frames_from_video(video_path: str, max_frames: Optional[int] = None) -> tuple[List[np.ndarray], float]:
"""
Extract frames from video file.
Args:
video_path (str): Path to video file.
max_frames (int, optional): Maximum number of frames to extract.
Returns:
tuple: (List of numpy arrays (H,W,3) uint8 RGB, FPS)
"""
if not os.path.exists(video_path):
raise FileNotFoundError(f"Video file not found: {video_path}")
cap = cv2.VideoCapture(video_path)
original_fps = cap.get(cv2.CAP_PROP_FPS)
all_frames = []
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Convert BGR to RGB
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
all_frames.append(frame_rgb)
cap.release()
if max_frames and len(all_frames) > max_frames:
frames = all_frames[:max_frames]
else:
frames = all_frames
return frames, original_fps
def run_inference(
pipeline: VideoInferencePipeline,
input_frames: List[np.ndarray],
mask_frames: List[np.ndarray],
chunk_size: int = 24 # Adjusted default chunk size
) -> List[np.ndarray]:
"""
Run VideoMaMa inference on video frames with mask conditioning.
Args:
pipeline (VideoInferencePipeline): Loaded pipeline instance.
input_frames (List[np.ndarray]): List of RGB frames (H,W,3) uint8.
mask_frames (List[np.ndarray]): List of mask frames (H,W) uint8 (0-255) grayscale.
chunk_size (int): Number of frames to process at once to avoid OOM.
Returns:
List[np.ndarray]: List of output RGB frames (H,W,3) uint8.
"""
if len(input_frames) != len(mask_frames):
# Resize mask frames list to match input if needed (e.g. repeat or slice)
# For strict correctness, we'll raise an error or warn.
# But let's assume the user provides matching lengths or we might need to handle it.
# Here we just raise for clarity.
raise ValueError(f"Input frames ({len(input_frames)}) and mask frames ({len(mask_frames)}) must have same length.")
# Convert numpy arrays to PIL Images
frames_pil = [Image.fromarray(f) for f in input_frames]
# Handle mask frames - ensure they are PIL "L" mode
mask_frames_pil = []
for m in mask_frames:
if m.ndim == 3:
# If RGB/BGR mask, convert to grayscale
m = cv2.cvtColor(m, cv2.COLOR_RGB2GRAY)
mask_frames_pil.append(Image.fromarray(m, mode='L'))
# Resize to model input size (1024x576 is standard for SVD)
target_width, target_height = 1024, 576
frames_resized = [f.resize((target_width, target_height), Image.Resampling.BILINEAR)
for f in frames_pil]
masks_resized = [m.resize((target_width, target_height), Image.Resampling.BILINEAR)
for m in mask_frames_pil]
print(f"Processing {len(frames_resized)} frames in chunks of {chunk_size}...")
# Store original size for resizing back
if not frames_pil:
return []
original_size = frames_pil[0].size
for i in range(0, len(frames_resized), chunk_size):
chunk_frames = frames_resized[i:i + chunk_size]
chunk_masks = masks_resized[i:i + chunk_size]
print(f" Running inference on chunk {i//chunk_size + 1}/{len(frames_resized)//chunk_size + 1} ({len(chunk_frames)} frames)...")
# Clear cache before each chunk
if pipeline.device.type == "cuda":
torch.cuda.empty_cache()
chunk_output = pipeline.run(
cond_frames=chunk_frames,
mask_frames=chunk_masks,
seed=42, # Fixed seed for reproducibility
mask_cond_mode="vae"
)
# Resize back to original resolution immediately
chunk_output_resized = [f.resize(original_size, Image.Resampling.BILINEAR)
for f in chunk_output]
# Convert back to numpy arrays
chunk_output_np = [np.array(f) for f in chunk_output_resized]
yield chunk_output_np
def save_video(frames: List[np.ndarray], output_path: str, fps: float):
"""
Save frames as a video file.
Args:
frames (List[np.ndarray]): List of frames (RGB).
output_path (str): Output video path.
fps (float): Frames per second.
"""
if not frames:
return
height, width = frames[0].shape[:2]
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
for frame in frames:
# Convert RGB to BGR for OpenCV
frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
out.write(frame_bgr)
out.release()
print(f"Saved video to {output_path}")
================================================
FILE: VideoMaMaInferenceModule/pipeline.py
================================================
# pipeline_svd_masked.py
import inspect
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Union
import numpy as np
import PIL.Image
import torch
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
from diffusers.image_processor import PipelineImageInput
from diffusers.models import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel
from diffusers.schedulers import EulerDiscreteScheduler
from diffusers.utils import BaseOutput, logging, replace_example_docstring
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
# Import necessary helpers from the original SVD pipeline
from diffusers.pipelines.stable_video_diffusion.pipeline_stable_video_diffusion import (
_append_dims,
retrieve_timesteps,
_resize_with_antialiasing,
)
import torch.nn.functional as F
from einops import rearrange
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from pipeline_svd_masked import StableVideoDiffusionPipelineWithMask
>>> from diffusers.utils import load_image, export_to_video
>>> # Load your fine-tuned UNet, VAE, etc.
>>> pipe = StableVideoDiffusionPipelineWithMask.from_pretrained(
... "path/to/your/finetuned_model", torch_dtype=torch.float16, variant="fp16"
... )
>>> pipe.to("cuda")
>>> # Load the conditioning image and the mask
>>> image = load_image("path/to/your/conditioning_image.png").resize((1024, 576))
>>> mask = load_image("path/to/your/mask_image.png").resize((1024, 576))
>>> # Generate frames
>>> frames = pipe(
... image=image,
... mask_image=mask,
... num_frames=25,
... decode_chunk_size=8
... ).frames[0]
>>> export_to_video(frames, "generated_video.mp4", fps=7)
```
"""
@dataclass
class StableVideoDiffusionPipelineOutput(BaseOutput):
r"""
Output class for the custom Stable Video Diffusion pipeline.
Args:
frames (`[List[List[PIL.Image.Image]]`, `np.ndarray`, `torch.Tensor`]):
List of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape
`(batch_size, num_frames, height, width, num_channels)`.
"""
frames: Union[List[List[PIL.Image.Image]], np.ndarray, torch.Tensor]
class StableVideoDiffusionPipelineWithMask(DiffusionPipeline):
r"""
A custom pipeline based on Stable Video Diffusion that accepts an additional mask for conditioning.
This pipeline is designed to work with a UNet fine-tuned to accept 12 input channels
(4 for noise, 4 for VAE-encoded condition image, 4 for VAE-encoded mask).
"""
model_cpu_offload_seq = "image_encoder->unet->vae"
_callback_tensor_inputs = ["latents"]
def __init__(
self,
vae: AutoencoderKLTemporalDecoder,
image_encoder: CLIPVisionModelWithProjection,
unet: UNetSpatioTemporalConditionModel,
scheduler: EulerDiscreteScheduler,
feature_extractor: CLIPImageProcessor,
):
super().__init__()
self.register_modules(
vae=vae,
image_encoder=image_encoder,
unet=unet,
scheduler=scheduler,
feature_extractor=feature_extractor,
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
self.video_processor = VideoProcessor(do_resize=True, vae_scale_factor=self.vae_scale_factor)
def _encode_image(
self,
image: PipelineImageInput,
device: Union[str, torch.device],
num_videos_per_prompt: int,
) -> torch.Tensor:
dtype = next(self.image_encoder.parameters()).dtype
if not isinstance(image, torch.Tensor):
image = self.video_processor.pil_to_numpy(image)
image = self.video_processor.numpy_to_pt(image)
image = image * 2.0 - 1.0
image = _resize_with_antialiasing(image, (224, 224))
image = (image + 1.0) / 2.0
image = self.feature_extractor(
images=image,
do_normalize=True,
do_center_crop=False,
do_resize=False,
do_rescale=False,
return_tensors="pt",
).pixel_values
image = image.to(device=device, dtype=dtype)
image_embeddings = self.image_encoder(image).image_embeds
image_embeddings = image_embeddings.unsqueeze(1)
bs_embed, seq_len, _ = image_embeddings.shape
image_embeddings = torch.zeros_like(image_embeddings)
return image_embeddings
def _encode_vae_image(
self,
image: torch.Tensor,
device: Union[str, torch.device],
num_videos_per_prompt: int,
):
image = image.to(device=device, dtype=torch.float16)
image_latents = self.vae.encode(image).latent_dist.sample()
image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1)
return image_latents
def _get_add_time_ids(
self,
fps: int,
motion_bucket_id: int,
noise_aug_strength: float,
dtype: torch.dtype,
batch_size: int,
num_videos_per_prompt: int,
):
add_time_ids = [fps, motion_bucket_id, noise_aug_strength]
passed_add_embed_dim = self.unet.config.addition_time_embed_dim * len(add_time_ids)
expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
if expected_add_embed_dim != passed_add_embed_dim:
raise ValueError(
f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created."
)
add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
add_time_ids = add_time_ids.repeat(batch_size * num_videos_per_prompt, 1)
return add_time_ids
def decode_latents(self, latents: torch.Tensor, num_frames: int, decode_chunk_size: int = 14):
latents = latents.flatten(0, 1).to(dtype=torch.float16)
latents = 1 / self.vae.config.scaling_factor * latents
frames = []
for i in range(0, latents.shape[0], decode_chunk_size):
num_frames_in = latents[i: i + decode_chunk_size].shape[0]
frame = self.vae.decode(latents[i: i + decode_chunk_size], num_frames=num_frames_in).sample
frames.append(frame)
frames = torch.cat(frames, dim=0)
frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4)
frames = frames.float()
return frames
def check_inputs(self, image, height, width):
if (
not isinstance(image, torch.Tensor)
and not isinstance(image, PIL.Image.Image)
and not isinstance(image, list)
):
raise ValueError(f"`image` has to be of type `torch.Tensor` or `PIL.Image.Image` but is {type(image)}")
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
def prepare_latents(
self,
batch_size: int,
num_frames: int,
height: int,
width: int,
dtype: torch.dtype,
device: Union[str, torch.device],
generator: torch.Generator,
latents: Optional[torch.Tensor] = None,
initial_latents: Optional[torch.Tensor] = None,
denoising_strength: float = 1.0,
timestep: Optional[torch.Tensor] = None,
):
num_channels_latents = self.unet.config.out_channels
shape = (
batch_size,
num_frames,
num_channels_latents,
height // self.vae_scale_factor,
width // self.vae_scale_factor,
)
if initial_latents is not None:
# Noise is added to the initial latents
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# Get the initial latents at the given timestep
latents = self.scheduler.add_noise(initial_latents, noise, timestep)
else:
# Standard pure noise generation
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
latents = latents.to(device)
# Scale the initial noise by the standard deviation required by the scheduler
latents = latents * self.scheduler.init_noise_sigma
return latents
def _encode_video_vae(
self,
video_frames: torch.Tensor, # Expects (B, F, C, H, W)
device: Union[str, torch.device],
):
video_frames = video_frames.to(device=device, dtype=self.vae.dtype)
batch_size, num_frames = video_frames.shape[:2]
# Reshape for VAE encoding
video_frames_reshaped = video_frames.reshape(batch_size * num_frames, *video_frames.shape[2:]) # (B*F, C, H, W)
latents = self.vae.encode(video_frames_reshaped).latent_dist.sample() # (B*F, C_latent, H_latent, W_latent)
# Reshape back to video format
latents = latents.reshape(batch_size, num_frames, *latents.shape[1:]) # (B, F, C_latent, H_latent, W_latent)
return latents
@torch.no_grad()
def __call__(
self,
image: Union[List[PIL.Image.Image], torch.Tensor],
mask_image: Union[List[PIL.Image.Image], torch.Tensor],
alpha_matte_image: Optional[Union[List[PIL.Image.Image], torch.Tensor]] = None,
denoising_strength: float = 0.7,
height: int = 576,
width: int = 1024,
num_frames: Optional[int] = None,
num_inference_steps: int = 30,
sigmas: Optional[List[float]] = None,
fps: int = 7,
motion_bucket_id: int = 127,
noise_aug_strength: float = 0.02,
decode_chunk_size: Optional[int] = None,
num_videos_per_prompt: Optional[int] = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
mask_noise_strength: float = 0.0,
):
height = height or self.unet.config.sample_size * self.vae_scale_factor
width = width or self.unet.config.sample_size * self.vae_scale_factor
if num_frames is None:
if isinstance(image, list):
num_frames = len(image)
else:
num_frames = self.unet.config.num_frames
decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else num_frames
self.check_inputs(image, height, width)
self.check_inputs(mask_image, height, width)
if alpha_matte_image:
self.check_inputs(alpha_matte_image, height, width)
batch_size = 1
device = self._execution_device
dtype = self.unet.dtype
image_for_clip = image[0] if isinstance(image, list) else image[0]
image_embeddings = self._encode_image(image_for_clip, device, num_videos_per_prompt)
fps = fps - 1
image_tensor = self.video_processor.preprocess(image, height=height, width=width).to(device).unsqueeze(0)
mask_tensor = self.video_processor.preprocess(mask_image, height=height, width=width).to(device).unsqueeze(0)
noise = randn_tensor(image_tensor.shape, generator=generator, device=device, dtype=dtype)
image_tensor = image_tensor + noise_aug_strength * noise
conditional_latents = self._encode_video_vae(image_tensor, device)
conditional_latents = conditional_latents / self.vae.config.scaling_factor
if self.unet.config.in_channels == 12:
mask_latents = self._encode_video_vae(mask_tensor, device)
mask_latents = mask_latents / self.vae.config.scaling_factor
elif self.unet.config.in_channels == 9:
mask_tensor_gray = mask_tensor.mean(dim=2, keepdim=True)
binarized_mask = (mask_tensor_gray > 0.0).to(dtype)
b, f, c, h, w = binarized_mask.shape
binarized_mask_reshaped = binarized_mask.reshape(b * f, c, h, w)
target_size = (height // self.vae_scale_factor, width // self.vae_scale_factor)
interpolated_mask = F.interpolate(
binarized_mask_reshaped,
size=target_size,
mode='nearest',
)
mask_latents = interpolated_mask.reshape(b, f, *interpolated_mask.shape[1:])
else:
raise ValueError(f"Unsupported number of UNet input channels: {self.unet.config.in_channels}.")
if mask_noise_strength > 0.0:
mask_noise = randn_tensor(mask_latents.shape, generator=generator, device=device, dtype=dtype)
mask_latents = mask_latents + mask_noise_strength * mask_noise
added_time_ids = self._get_add_time_ids(
fps, motion_bucket_id, noise_aug_strength, image_embeddings.dtype, batch_size, num_videos_per_prompt
)
added_time_ids = added_time_ids.to(device)
# --- MODIFIED FOR ALPHA MATTE REFINEMENT ---
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, None, sigmas)
# self.scheduler.set_timesteps(num_inference_steps, device=device)
# timesteps = self.scheduler.timesteps
initial_latents = None
if alpha_matte_image is not None:
alpha_matte_tensor = self.video_processor.preprocess(alpha_matte_image, height=height, width=width).to(
device).unsqueeze(0)
initial_latents = self._encode_video_vae(alpha_matte_tensor, device)
initial_latents = initial_latents / self.vae.config.scaling_factor
# Adjust the number of steps and the timesteps to start from
t_start = max(num_inference_steps - int(num_inference_steps * denoising_strength), 0)
timesteps = timesteps[t_start:]
# We need the first timestep to add the correct amount of noise
start_timestep = timesteps[0]
else:
start_timestep = timesteps[0] # Not used, but for clarity
latents = self.prepare_latents(
batch_size * num_videos_per_prompt,
num_frames,
height,
width,
dtype,
device,
generator,
latents,
initial_latents=initial_latents,
denoising_strength=denoising_strength,
timestep=start_timestep if initial_latents is not None else None,
)
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
with self.progress_bar(total=len(timesteps)) as progress_bar:
for i, t in enumerate(timesteps):
latent_model_input = self.scheduler.scale_model_input(latents, t)
latent_model_input = torch.cat([latent_model_input, conditional_latents, mask_latents], dim=2)
noise_pred = self.unet(
latent_model_input, t, encoder_hidden_states=image_embeddings, added_time_ids=added_time_ids,
return_dict=False
)[0]
latents = self.scheduler.step(noise_pred, t, latents).prev_sample
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
frames = self.decode_latents(latents, num_frames, decode_chunk_size)
frames = self.video_processor.postprocess_video(video=frames, output_type=output_type)
self.maybe_free_model_hooks()
if not return_dict:
return frames
return StableVideoDiffusionPipelineOutput(frames=frames)
class StableVideoDiffusionPipelineOnestepWithMask(DiffusionPipeline):
r"""
A custom pipeline based on Stable Video Diffusion that accepts an additional mask for conditioning.
This pipeline is designed to work with a UNet fine-tuned to accept 12 input channels
(4 for noise, 4 for VAE-encoded condition image, 4 for VAE-encoded mask).
"""
model_cpu_offload_seq = "image_encoder->unet->vae"
_callback_tensor_inputs = ["latents"]
def __init__(
self,
vae: AutoencoderKLTemporalDecoder,
image_encoder: CLIPVisionModelWithProjection,
unet: UNetSpatioTemporalConditionModel,
scheduler: EulerDiscreteScheduler,
feature_extractor: CLIPImageProcessor,
):
super().__init__()
self.register_modules(
vae=vae,
image_encoder=image_encoder,
unet=unet,
scheduler=scheduler,
feature_extractor=feature_extractor,
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
self.video_processor = VideoProcessor(do_resize=True, vae_scale_factor=self.vae_scale_factor)
def _encode_image(
self,
image: PipelineImageInput,
device: Union[str, torch.device],
num_videos_per_prompt: int,
) -> torch.Tensor:
dtype = next(self.image_encoder.parameters()).dtype
if not isinstance(image, torch.Tensor):
image = self.video_processor.pil_to_numpy(image)
image = self.video_processor.numpy_to_pt(image)
image = image * 2.0 - 1.0
image = _resize_with_antialiasing(image, (224, 224))
image = (image + 1.0) / 2.0
image = self.feature_extractor(
images=image,
do_normalize=True,
do_center_crop=False,
do_resize=False,
do_rescale=False,
return_tensors="pt",
).pixel_values
image = image.to(device=device, dtype=dtype)
image_embeddings = self.image_encoder(image).image_embeds
image_embeddings = image_embeddings.unsqueeze(1)
bs_embed, seq_len, _ = image_embeddings.shape
image_embeddings = torch.zeros_like(image_embeddings)
return image_embeddings
def _encode_vae_image(
self,
image: torch.Tensor,
device: Union[str, torch.device],
num_videos_per_prompt: int,
):
image = image.to(device=device, dtype=torch.float16)
image_latents = self.vae.encode(image).latent_dist.sample()
image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1)
return image_latents
def _get_add_time_ids(
self,
fps: int,
motion_bucket_id: int,
noise_aug_strength: float,
dtype: torch.dtype,
batch_size: int,
num_videos_per_prompt: int,
):
add_time_ids = [fps, motion_bucket_id, noise_aug_strength]
passed_add_embed_dim = self.unet.config.addition_time_embed_dim * len(add_time_ids)
expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
if expected_add_embed_dim != passed_add_embed_dim:
raise ValueError(
f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created."
)
add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
add_time_ids = add_time_ids.repeat(batch_size * num_videos_per_prompt, 1)
return add_time_ids
def decode_latents(self, latents: torch.Tensor, num_frames: int, decode_chunk_size: int = 14):
latents = latents.flatten(0, 1).to(dtype=torch.float16)
latents = 1 / self.vae.config.scaling_factor * latents
frames = []
for i in range(0, latents.shape[0], decode_chunk_size):
num_frames_in = latents[i: i + decode_chunk_size].shape[0]
frame = self.vae.decode(latents[i: i + decode_chunk_size], num_frames=num_frames_in).sample
frames.append(frame)
frames = torch.cat(frames, dim=0)
frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4)
frames = frames.float()
return frames
def check_inputs(self, image, height, width):
if (
not isinstance(image, torch.Tensor)
and not isinstance(image, PIL.Image.Image)
and not isinstance(image, list)
):
raise ValueError(f"`image` has to be of type `torch.Tensor` or `PIL.Image.Image` but is {type(image)}")
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
def prepare_latents(
self,
batch_size: int,
num_frames: int,
height: int,
width: int,
dtype: torch.dtype,
device: Union[str, torch.device],
generator: torch.Generator,
latents: Optional[torch.Tensor] = None,
):
# The number of channels for the initial noise is based on the UNet's out_channels
num_channels_latents = self.unet.config.out_channels
shape = (
batch_size,
num_frames,
num_channels_latents,
height // self.vae_scale_factor,
width // self.vae_scale_factor,
)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(f"batch size {batch_size} must match the length of the generators {len(generator)}.")
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
latents = latents.to(device)
latents = latents * self.scheduler.init_noise_sigma
return latents
def _encode_video_vae(
self,
video_frames: torch.Tensor, # Expects (B, F, C, H, W)
device: Union[str, torch.device],
):
video_frames = video_frames.to(device=device, dtype=self.vae.dtype)
batch_size, num_frames = video_frames.shape[:2]
# Reshape for VAE encoding
video_frames_reshaped = video_frames.reshape(batch_size * num_frames, *video_frames.shape[2:]) # (B*F, C, H, W)
latents = self.vae.encode(video_frames_reshaped).latent_dist.sample() # (B*F, C_latent, H_latent, W_latent)
# Reshape back to video format
latents = latents.reshape(batch_size, num_frames, *latents.shape[1:]) # (B, F, C_latent, H_latent, W_latent)
return latents
@torch.no_grad()
def __call__(
self,
image: Union[List[PIL.Image.Image], torch.Tensor],
mask_image: Union[List[PIL.Image.Image], torch.Tensor],
height: int = 576,
width: int = 1024,
num_frames: Optional[int] = None,
fps: int = 7,
motion_bucket_id: int = 127,
noise_aug_strength: float = 0.0,
decode_chunk_size: Optional[int] = None,
num_videos_per_prompt: Optional[int] = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
mask_noise_strength: float = 0.0,
):
height = height or self.unet.config.sample_size * self.vae_scale_factor
width = width or self.unet.config.sample_size * self.vae_scale_factor
if num_frames is None:
if isinstance(image, list):
num_frames = len(image)
else:
num_frames = self.unet.config.num_frames
decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else num_frames
self.check_inputs(image, height, width)
self.check_inputs(mask_image, height, width)
if isinstance(image, list) and isinstance(mask_image, list):
if len(image) != len(mask_image):
raise ValueError("`image` and `mask_image` must have the same number of frames.")
if num_frames != len(image):
logger.warning(
f"Mismatch between `num_frames` ({num_frames}) and number of input images ({len(image)}). Using {len(image)}.")
num_frames = len(image)
batch_size = 1
device = self._execution_device
dtype = self.unet.dtype
image_for_clip = image[0] if isinstance(image, list) else image[0]
image_embeddings = self._encode_image(image_for_clip, device, num_videos_per_prompt)
fps = fps - 1
image_tensor = self.video_processor.preprocess(image, height=height, width=width).to(device).unsqueeze(0)
mask_tensor = self.video_processor.preprocess(mask_image, height=height, width=width).to(
device).unsqueeze(0)
noise = randn_tensor(image_tensor.shape, generator=generator, device=device, dtype=dtype)
image_tensor = image_tensor + noise_aug_strength * noise
conditional_latents = self._encode_video_vae(image_tensor, device)
conditional_latents = conditional_latents / self.vae.config.scaling_factor
if self.unet.config.in_channels == 12:
mask_latents = self._encode_video_vae(mask_tensor, device)
mask_latents = mask_latents / self.vae.config.scaling_factor
elif self.unet.config.in_channels == 9:
mask_tensor_gray = mask_tensor.mean(dim=2, keepdim=True)
binarized_mask = (mask_tensor_gray > 0.0).to(dtype)
b, f, c, h, w = binarized_mask.shape
binarized_mask_reshaped = binarized_mask.reshape(b * f, c, h, w)
target_size = (height // self.vae_scale_factor, width // self.vae_scale_factor)
interpolated_mask = F.interpolate(
binarized_mask_reshaped,
size=target_size,
mode='nearest',
)
mask_latents = interpolated_mask.reshape(b, f, *interpolated_mask.shape[1:])
else:
raise ValueError(
f"Unsupported number of UNet input channels: {self.unet.config.in_channels}. "
"This pipeline only supports 9 (for interpolated mask) or 12 (for VAE mask)."
)
if mask_noise_strength > 0.0:
mask_noise = randn_tensor(mask_latents.shape, generator=generator, device=device, dtype=dtype)
mask_latents = mask_latents + mask_noise_strength * mask_noise
added_time_ids = self._get_add_time_ids(
fps, motion_bucket_id, noise_aug_strength, image_embeddings.dtype, batch_size, num_videos_per_prompt
)
added_time_ids = added_time_ids.to(device)
# **MODIFIED FOR SINGLE-STEP**: Prepare initial noise
num_channels_latents = self.unet.config.out_channels
shape = (
batch_size * num_videos_per_prompt,
num_frames,
num_channels_latents,
height // self.vae_scale_factor,
width // self.vae_scale_factor,
)
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# **MODIFIED FOR SINGLE-STEP**: Set a fixed high timestep
timestep = torch.tensor([1.0], dtype=dtype, device=device) # Use a high sigma value
# **MODIFIED FOR SINGLE-STEP**: Single forward pass
latent_model_input = torch.cat([latents, conditional_latents, mask_latents], dim=2)
noise_pred = self.unet(
latent_model_input, timestep, encoder_hidden_states=image_embeddings, added_time_ids=added_time_ids,
return_dict=False
)[0]
# The model's prediction is the final denoised latent
denoised_latents = noise_pred
frames = self.decode_latents(denoised_latents, num_frames, decode_chunk_size)
frames = self.video_processor.postprocess_video(video=frames, output_type=output_type)
self.maybe_free_model_hooks()
if not return_dict:
return frames
return StableVideoDiffusionPipelineOutput(frames=frames)
class StableVideoDiffusionPipelineWithCrossAtnnMask(DiffusionPipeline):
model_cpu_offload_seq = "image_encoder->unet->vae"
_callback_tensor_inputs = ["latents"]
def __init__(
self,
vae: AutoencoderKLTemporalDecoder,
unet: UNetSpatioTemporalConditionModel,
scheduler: EulerDiscreteScheduler,
mask_projector: torch.nn.Module,
# CLIP models are not strictly needed for inference if embeddings are not used
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
):
super().__init__()
self.register_modules(
vae=vae,
unet=unet,
scheduler=scheduler,
mask_projector=mask_projector,
image_encoder=image_encoder,
feature_extractor=feature_extractor,
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
self.video_processor = VideoProcessor(do_resize=False, vae_scale_factor=self.vae_scale_factor)
def _encode_image_vae(self, image: torch.Tensor, device: Union[str, torch.device]):
image = image.to(device=device, dtype=self.vae.dtype)
latent = self.vae.encode(image).latent_dist.sample()
return latent
def decode_latents(self, latents: torch.Tensor, num_frames: int, decode_chunk_size: int):
latents = latents.flatten(0, 1).to(dtype=torch.float16)
latents = 1 / self.vae.config.scaling_factor * latents
frames = []
for i in range(0, latents.shape[0], decode_chunk_size):
frame = self.vae.decode(latents[i: i + decode_chunk_size], num_frames=decode_chunk_size).sample
frames.append(frame)
frames = torch.cat(frames, dim=0)
frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4)
frames = frames.float()
return frames
def _encode_video_vae(
self,
video_frames: torch.Tensor, # Expects (B, F, C, H, W)
device: Union[str, torch.device],
):
video_frames = video_frames.to(device=device, dtype=self.vae.dtype)
batch_size, num_frames = video_frames.shape[:2]
# Reshape for VAE encoding
video_frames_reshaped = video_frames.reshape(batch_size * num_frames, *video_frames.shape[2:]) # (B*F, C, H, W)
latents = self.vae.encode(video_frames_reshaped).latent_dist.sample() # (B*F, C_latent, H_latent, W_latent)
# Reshape back to video format
latents = latents.reshape(batch_size, num_frames, *latents.shape[1:]) # (B, F, C_latent, H_latent, W_latent)
return latents
@torch.no_grad()
def __call__(
self,
image: Union[PIL.Image.Image, torch.Tensor], # Static image for appearance
mask_image: List[PIL.Image.Image], # Video mask for motion
height: int = 576,
width: int = 1024,
num_frames: Optional[int] = None,
num_inference_steps: int = 25,
fps: int = 7,
motion_bucket_id: int = 127,
noise_aug_strength: float = 0.0, # Noise is added to latents now
decode_chunk_size: Optional[int] = 8,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
):
device = self._execution_device
dtype = self.unet.dtype
num_frames = num_frames if num_frames is not None else len(mask_image)
decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else num_frames
# 1. PREPARE STATIC IMAGE CONDITION
image_tensor = self.video_processor.preprocess(image, height, width).to(device).unsqueeze(0)
conditional_latents = self._encode_video_vae(image_tensor, device)
conditional_latents = conditional_latents / self.vae.config.scaling_factor
# 2. PREPARE MASK MOTION CONDITION
mask_tensor = self.video_processor.preprocess(mask_image, height, width)
if mask_tensor.shape[1] > 1:
mask_tensor = mask_tensor.mean(dim=1, keepdim=True)
# Reshape for projector: (T, C, H, W)
mask_for_projection = rearrange(mask_tensor, "f c h w -> f c h w").to(device, dtype)
encoder_hidden_states = self.mask_projector(mask_for_projection)
encoder_hidden_states = encoder_hidden_states.unsqueeze(1) # (T, 1, D)
# Add batch dimension for UNet
encoder_hidden_states = encoder_hidden_states.unsqueeze(0) # (1, T, 1, D)
# The UNet will handle flattening this to (B*T, 1, D) where B=1
# To be safe, we pass it pre-flattened.
encoder_hidden_states = rearrange(encoder_hidden_states, "b f s d -> (b f) s d")
# 3. PREPARE LATENTS
shape = (1, num_frames, self.unet.config.out_channels, height // self.vae_scale_factor,
width // self.vae_scale_factor)
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
if noise_aug_strength > 0:
latents += noise_aug_strength * randn_tensor(latents.shape, generator=generator, device=device,
dtype=dtype)
latents = latents * self.scheduler.init_noise_sigma
# 4. GET ADDED TIME IDS
# For pipeline, batch size is 1
added_time_ids = [fps - 1, motion_bucket_id, 0.0] # noise_aug_strength for add_time_ids is 0 for inference
added_time_ids = torch.tensor([added_time_ids], dtype=dtype, device=device)
# 5. DENOISING LOOP
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
with self.progress_bar(total=num_inference_steps) as progress_bar:
for t in timesteps:
latent_model_input = self.scheduler.scale_model_input(latents, t)
unet_input = torch.cat([latent_model_input, conditional_latents], dim=2)
noise_pred = self.unet(
unet_input, t, encoder_hidden_states=encoder_hidden_states, added_time_ids=added_time_ids
).sample
latents = self.scheduler.step(noise_pred, t, latents).prev_sample
progress_bar.update()
# 6. DECODE
frames = self.decode_latents(latents, num_frames, decode_chunk_size)
frames = self.video_processor.postprocess_video(video=frames, output_type=output_type)
if not return_dict:
return (frames,)
return StableVideoDiffusionPipelineOutput(frames=frames)
# pipeline.py
import torch
import torch.nn.functional as F
from PIL import Image
from einops import rearrange
from torchvision import transforms
from diffusers import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
class VideoInferencePipeline:
"""
A reusable pipeline for single-step video diffusion inference.
This class encapsulates the models and the core inference logic,
separating it from data loading and saving, which can vary between tasks.
"""
def __init__(self, base_model_path: str, unet_checkpoint_path: str, device: str = "cuda",
weight_dtype: torch.dtype = torch.float16):
"""
Loads all necessary models into memory.
Args:
base_model_path (str): Path to the base Stable Video Diffusion model.
unet_checkpoint_path (str): Path to the fine-tuned UNet checkpoint.
device (str): The device to run models on ('cuda' or 'cpu').
weight_dtype (torch.dtype): The precision for model weights (float16 or bfloat16).
"""
logger.info("--- Initializing Inference Pipeline and Loading Models ---")
self.device = torch.device(device if torch.cuda.is_available() else "cpu")
self.weight_dtype = weight_dtype
# Load models from pretrained paths
try:
self.feature_extractor = CLIPImageProcessor.from_pretrained(base_model_path, subfolder="feature_extractor")
self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(base_model_path,
subfolder="image_encoder",
variant="fp16")
self.vae = AutoencoderKLTemporalDecoder.from_pretrained(base_model_path, subfolder="vae", variant="fp16")
self.unet = UNetSpatioTemporalConditionModel.from_pretrained(unet_checkpoint_path, subfolder="unet")
except Exception as e:
raise IOError(f"Fatal error loading models: {e}")
# Move models to the specified device and set to evaluation mode
# CLIP must run in FP32 to avoid CUBLAS errors
self.image_encoder.to(self.device, dtype=torch.float32).eval()
# VAE must also run in FP32 to avoid CUBLAS errors
self.vae.to(self.device, dtype=torch.float32).eval()
self.unet.to(self.device, dtype=self.weight_dtype).eval()
logger.info("--- Models Loaded Successfully on %s ---", self.device)
def run(self, cond_frames, mask_frames, seed=42, mask_cond_mode="vae", fps=7, motion_bucket_id=127,
noise_aug_strength=0.0):
"""
Runs the core inference process on a sequence of conditioning and mask frames.
Args:
cond_frames (list[Image.Image]): List of PIL images for conditioning.
mask_frames (list[Image.Image]): List of PIL images for the masks.
seed (int): Random seed for generation.
mask_cond_mode (str): How the mask is conditioned ("vae" or "interpolate").
fps (int): Frames per second to condition the model with.
motion_bucket_id (int): Motion bucket ID for conditioning.
noise_aug_strength (float): Noise augmentation strength.
Returns:
list[Image.Image]: A list of the generated video frames as PIL Images.
"""
# --- 1. Prepare Tensors ---
cond_video_tensor = self._pil_to_tensor(cond_frames).to(self.device)
mask_video_tensor = self._pil_to_tensor(mask_frames).to(self.device)
if mask_video_tensor.shape[2] != 3:
mask_video_tensor = mask_video_tensor.repeat(1, 1, 3, 1, 1)
with torch.no_grad():
# --- 2. Get CLIP Image Embeddings ---
first_frame_tensor = cond_video_tensor[:, 0, :, :, :]
pixel_values_for_clip = self._resize_with_antialiasing(first_frame_tensor, (224, 224))
pixel_values_for_clip = ((pixel_values_for_clip + 1.0) / 2.0).clamp(0, 1)
pixel_values = self.feature_extractor(images=pixel_values_for_clip, do_rescale=False, return_tensors="pt").pixel_values
# Run CLIP in FP32
image_embeddings = self.image_encoder(pixel_values.to(self.device, dtype=torch.float32)).image_embeds
logger.debug("CLIP Embeds Max: %.4f, Mean: %.4f", image_embeddings.max().item(), image_embeddings.mean().item())
# Setup for UNet which uses weight_dtype (likely FP16)
image_embeddings = image_embeddings.to(dtype=self.weight_dtype)
encoder_hidden_states = torch.zeros_like(image_embeddings).unsqueeze(1)
# --- 3. Prepare Latents ---
# VAE encoding must happen in FP32
cond_video_tensor_fp32 = cond_video_tensor.to(dtype=torch.float32)
cond_latents = self._tensor_to_vae_latent(cond_video_tensor_fp32)
logger.debug("Cond Latents Max: %.4f, Mean: %.4f", cond_latents.max().item(), cond_latents.mean().item())
# Cast back to weight_dtype (FP16) for UNet
cond_latents = cond_latents.to(dtype=self.weight_dtype)
cond_latents = cond_latents / self.vae.config.scaling_factor
if mask_cond_mode == "vae":
mask_video_tensor_fp32 = mask_video_tensor.to(dtype=torch.float32)
mask_latents = self._tensor_to_vae_latent(mask_video_tensor_fp32)
logger.debug("Mask Latents Max: %.4f, Mean: %.4f", mask_latents.max().item(), mask_latents.mean().item())
mask_latents = mask_latents.to(dtype=self.weight_dtype)
mask_latents = mask_latents / self.vae.config.scaling_factor
elif mask_cond_mode == "interpolate":
target_shape = cond_latents.shape[-2:]
b, t, c, h, w = mask_video_tensor.shape
mask_video_reshaped = rearrange(mask_video_tensor, "b t c h w -> (b t) c h w")
interpolated_mask = F.interpolate(mask_video_reshaped, size=target_shape, mode='bilinear',
align_corners=False)
mask_latents = rearrange(interpolated_mask, "(b t) c h w -> b t c h w", b=b)
else:
raise ValueError(f"Unknown mask_cond_mode: {mask_cond_mode}")
# --- 4. Run UNet Single-Step Inference ---
generator = torch.Generator(device=self.device).manual_seed(seed)
noisy_latents = torch.randn(cond_latents.shape, generator=generator, device=self.device,
dtype=self.weight_dtype)
timesteps = torch.full((1,), 1.0, device=self.device, dtype=torch.long)
added_time_ids = self._get_add_time_ids(fps, motion_bucket_id, noise_aug_strength, batch_size=1)
unet_input = torch.cat([noisy_latents, cond_latents, mask_latents], dim=2)
pred_latents = self.unet(unet_input, timesteps, encoder_hidden_states, added_time_ids=added_time_ids).sample
logger.debug("Pred Latents Max: %.4f, Mean: %.4f", pred_latents.max().item(), pred_latents.mean().item())
# --- 5. Decode Latents to Video Frames ---
pred_latents = (1 / self.vae.config.scaling_factor) * pred_latents.squeeze(0)
frames = []
# Process in chunks to avoid VRAM issues, especially for long videos
# Decode in FP32
pred_latents_fp32 = pred_latents.to(dtype=torch.float32)
for i in range(0, pred_latents_fp32.shape[0], 8):
chunk = pred_latents_fp32[i: i + 8]
decoded_chunk = self.vae.decode(chunk, num_frames=chunk.shape[0]).sample
frames.append(decoded_chunk)
video_tensor = torch.cat(frames, dim=0)
logger.debug("Video Tensor (Pre-Clamp) Max: %.4f, Mean: %.4f", video_tensor.max().item(), video_tensor.mean().item())
video_tensor = (video_tensor / 2.0 + 0.5).clamp(0, 1).mean(dim=1, keepdim=True).repeat(1, 3, 1, 1)
# Return a list of PIL images
return [transforms.ToPILImage()(frame) for frame in video_tensor]
def _pil_to_tensor(self, frames: list[Image.Image]):
"""Converts a list of PIL images to a normalized video tensor."""
video_tensor = torch.stack([transforms.ToTensor()(f) for f in frames]).unsqueeze(0)
return video_tensor * 2.0 - 1.0
def _tensor_to_vae_latent(self, t: torch.Tensor):
"""Encodes a video tensor into the VAE's latent space in chunks to avoid OOM."""
video_length = t.shape[1]
t = rearrange(t, "b f c h w -> (b f) c h w")
# Process in chunks of 8
chunk_size = 8
latents_list = []
for i in range(0, t.shape[0], chunk_size):
chunk = t[i:i + chunk_size]
chunk_latents = self.vae.encode(chunk).latent_dist.sample()
latents_list.append(chunk_latents)
latents = torch.cat(latents_list, dim=0)
latents = rearrange(latents, "(b f) c h w -> b f c h w", f=video_length)
return latents * self.vae.config.scaling_factor
def _get_add_time_ids(self, fps, motion_bucket_id, noise_aug_strength, batch_size):
"""Creates the additional time IDs for conditioning the UNet."""
add_time_ids_list = [fps, motion_bucket_id, noise_aug_strength]
passed_add_embed_dim = self.unet.config.addition_time_embed_dim * len(add_time_ids_list)
expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
if expected_add_embed_dim != passed_add_embed_dim:
raise ValueError(
f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created.")
add_time_ids = torch.tensor([add_time_ids_list], dtype=self.weight_dtype, device=self.device)
return add_time_ids.repeat(batch_size, 1)
def _resize_with_antialiasing(self, input_tensor, size, interpolation="bicubic", align_corners=True):
"""
Resizes a tensor with anti-aliasing for CLIP input, mirroring k-diffusion.
This is a direct copy of the helper function from your original scripts.
"""
h, w = input_tensor.shape[-2:]
factors = (h / size[0], w / size[1])
sigmas = (max((factors[0] - 1.0) / 2.0, 0.001), max((factors[1] - 1.0) / 2.0, 0.001))
ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3))
if (ks[0] % 2) == 0: ks = ks[0] + 1, ks[1]
if (ks[1] % 2) == 0: ks = ks[0], ks[1] + 1
def _compute_padding(kernel_size):
computed = [k - 1 for k in kernel_size]
out_padding = 2 * len(kernel_size) * [0]
for i in range(len(kernel_size)):
computed_tmp = computed[-(i + 1)]
pad_front = computed_tmp // 2
pad_rear = computed_tmp - pad_front
out_padding[2 * i + 0] = pad_front
out_padding[2 * i + 1] = pad_rear
return out_padding
def _filter2d(input_tensor, kernel):
b, c, h, w = input_tensor.shape
tmp_kernel = kernel[:, None, ...].to(device=input_tensor.device, dtype=input_tensor.dtype)
tmp_kernel = tmp_kernel.expand(-1, c, -1, -1)
height, width = tmp_kernel.shape[-2:]
padding_shape = _compute_padding([height, width])
input_tensor_padded = F.pad(input_tensor, padding_shape, mode="reflect")
tmp_kernel = tmp_kernel.reshape(-1, 1, height, width)
input_tensor_padded = input_tensor_padded.view(-1, tmp_kernel.size(0), input_tensor_padded.size(-2),
input_tensor_padded.size(-1))
output = F.conv2d(input_tensor_padded, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1)
return output.view(b, c, h, w)
def _gaussian(window_size, sigma):
if isinstance(sigma, float):
sigma = torch.tensor([[sigma]])
x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(
sigma.shape[0], -1)
if window_size % 2 == 0:
x = x + 0.5
gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0)))
return gauss / gauss.sum(-1, keepdim=True)
def _gaussian_blur2d(input_tensor, kernel_size, sigma):
if isinstance(sigma, tuple):
sigma = torch.tensor([sigma], dtype=input_tensor.dtype)
else:
sigma = sigma.to(dtype=input_tensor.dtype)
ky, kx = int(kernel_size[0]), int(kernel_size[1])
bs = sigma.shape[0]
kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1))
kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1))
out_x = _filter2d(input_tensor, kernel_x[..., None, :])
return _filter2d(out_x, kernel_y[..., None])
blurred_input = _gaussian_blur2d(input_tensor, ks, sigmas)
return F.interpolate(blurred_input, size=size, mode=interpolation, align_corners=align_corners)
================================================
FILE: backend/__init__.py
================================================
"""Backend service layer for ez-CorridorKey."""
from .clip_state import (
ClipAsset,
ClipEntry,
ClipState,
InOutRange,
scan_clips_dir,
scan_project_clips,
)
from .errors import CorridorKeyError
from .job_queue import GPUJob, GPUJobQueue, JobStatus, JobType
from .natural_sort import natsorted, natural_sort_key
from .project import (
VIDEO_FILE_FILTER,
add_clips_to_project,
create_project,
get_clip_dirs,
get_display_name,
is_image_file,
is_v2_project,
is_video_file,
projects_root,
read_clip_json,
read_project_json,
sanitize_stem,
set_display_name,
write_clip_json,
write_project_json,
)
from .service import CorridorKeyService, InferenceParams, OutputConfig
__all__ = [
# Service
"CorridorKeyService",
"InferenceParams",
"OutputConfig",
# Clip state
"ClipAsset",
"ClipEntry",
"ClipState",
"InOutRange",
"scan_clips_dir",
"scan_project_clips",
# Job queue
"GPUJob",
"GPUJobQueue",
"JobType",
"JobStatus",
# Errors
"CorridorKeyError",
# Project utilities
"projects_root",
"create_project",
"add_clips_to_project",
"sanitize_stem",
"get_clip_dirs",
"is_v2_project",
"write_project_json",
"read_project_json",
"write_clip_json",
"read_clip_json",
"get_display_name",
"set_display_name",
"is_video_file",
"is_image_file",
"VIDEO_FILE_FILTER",
# Natural sort
"natural_sort_key",
"natsorted",
]
================================================
FILE: backend/clip_state.py
================================================
"""Clip entry data model and state machine.
State Machine:
EXTRACTING — Video input being extracted to image sequence
RAW — Input asset found, no alpha hint yet
MASKED — User mask provided (for VideoMaMa workflow)
READY — Alpha hint available (from GVM or VideoMaMa), ready for inference
COMPLETE — Inference outputs written
ERROR — Processing failed (can retry)
Transitions:
EXTRACTING → RAW (extraction completes)
EXTRACTING → ERROR (extraction fails)
RAW → MASKED (user provides VideoMaMa mask)
RAW → READY (GVM auto-generates alpha)
RAW → ERROR (GVM/scan fails)
MASKED → READY (VideoMaMa generates alpha from user mask)
MASKED → ERROR (VideoMaMa fails)
READY → COMPLETE (inference succeeds)
READY → ERROR (inference fails)
ERROR → RAW (retry from scratch)
ERROR → MASKED (retry with mask)
ERROR → READY (retry inference)
ERROR → EXTRACTING (retry extraction)
COMPLETE → READY (reprocess with different params)
"""
from __future__ import annotations
import glob as glob_module
import logging
import os
from dataclasses import dataclass, field
from enum import Enum
from .errors import ClipScanError, InvalidStateTransitionError
from .natural_sort import natsorted
from .project import is_image_file as _is_image_file
from .project import is_video_file as _is_video_file
logger = logging.getLogger(__name__)
class ClipState(Enum):
EXTRACTING = "EXTRACTING"
RAW = "RAW"
MASKED = "MASKED"
READY = "READY"
COMPLETE = "COMPLETE"
ERROR = "ERROR"
# Valid transitions: from_state -> set of allowed to_states
_TRANSITIONS: dict[ClipState, set[ClipState]] = {
ClipState.EXTRACTING: {ClipState.RAW, ClipState.ERROR},
ClipState.RAW: {ClipState.MASKED, ClipState.READY, ClipState.ERROR},
ClipState.MASKED: {ClipState.READY, ClipState.ERROR},
ClipState.READY: {ClipState.COMPLETE, ClipState.ERROR},
ClipState.COMPLETE: {ClipState.READY}, # reprocess with different params
ClipState.ERROR: {ClipState.RAW, ClipState.MASKED, ClipState.READY, ClipState.EXTRACTING},
}
@dataclass
class ClipAsset:
"""Represents an input source — either an image sequence directory or a video file."""
path: str
asset_type: str # 'sequence' or 'video'
frame_count: int = 0
def __post_init__(self):
self._calculate_length()
def _calculate_length(self):
if self.asset_type == "sequence":
if os.path.isdir(self.path):
files = [f for f in os.listdir(self.path) if _is_image_file(f)]
self.frame_count = len(files)
else:
self.frame_count = 0
elif self.asset_type == "video":
try:
import cv2
cap = cv2.VideoCapture(self.path)
if cap.isOpened():
self.frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
if self.frame_count == 0:
logger.warning(f"Video reports 0 frames, file may be corrupted: {self.path}")
cap.release()
except Exception as e:
logger.debug(f"Video frame count detection failed for {self.path}: {e}")
self.frame_count = 0
def get_frame_files(self) -> list[str]:
"""Return naturally sorted list of frame filenames for sequence assets.
Uses natural sort so frame_2 sorts before frame_10 (not lexicographic).
"""
if self.asset_type != "sequence" or not os.path.isdir(self.path):
return []
return natsorted([f for f in os.listdir(self.path) if _is_image_file(f)])
@dataclass
class InOutRange:
"""In/out frame range for sub-clip processing. Both indices inclusive, 0-based."""
in_point: int
out_point: int
@property
def frame_count(self) -> int:
return self.out_point - self.in_point + 1
def contains(self, index: int) -> bool:
return self.in_point <= index <= self.out_point
def to_dict(self) -> dict:
return {"in_point": self.in_point, "out_point": self.out_point}
@classmethod
def from_dict(cls, d: dict) -> InOutRange:
return cls(in_point=d["in_point"], out_point=d["out_point"])
@dataclass
class ClipEntry:
"""A single shot/clip with its assets and processing state."""
name: str
root_path: str
state: ClipState = ClipState.RAW
input_asset: ClipAsset | None = None
alpha_asset: ClipAsset | None = None
mask_asset: ClipAsset | None = None # User-provided VideoMaMa mask
in_out_range: InOutRange | None = None # Per-clip in/out markers (None = full clip)
warnings: list[str] = field(default_factory=list)
error_message: str | None = None
extraction_progress: float = 0.0 # 0.0 to 1.0 during EXTRACTING
extraction_total: int = 0 # total frames expected during extraction
_processing: bool = field(default=False, repr=False) # lock: watcher must not reclassify
@property
def is_processing(self) -> bool:
"""True while a GPU job is actively working on this clip."""
return self._processing
def set_processing(self, value: bool) -> None:
"""Set processing lock. Watcher skips reclassification while True."""
self._processing = value
def transition_to(self, new_state: ClipState) -> None:
"""Attempt a state transition. Raises InvalidStateTransitionError if not allowed."""
if new_state not in _TRANSITIONS.get(self.state, set()):
raise InvalidStateTransitionError(self.name, self.state.value, new_state.value)
old = self.state
self.state = new_state
if new_state != ClipState.ERROR:
self.error_message = None
logger.debug(f"Clip '{self.name}': {old.value} -> {new_state.value}")
def set_error(self, message: str) -> None:
"""Transition to ERROR state with a message.
Works from any state that allows ERROR transition
(RAW, MASKED, READY — all can error now).
"""
self.transition_to(ClipState.ERROR)
self.error_message = message
@property
def output_dir(self) -> str:
return os.path.join(self.root_path, "Output")
@property
def has_outputs(self) -> bool:
"""Check if output directory exists with content."""
out = self.output_dir
if not os.path.isdir(out):
return False
for subdir in ("FG", "Matte", "Comp", "Processed"):
d = os.path.join(out, subdir)
if os.path.isdir(d) and os.listdir(d):
return True
return False
def completed_frame_count(self) -> int:
"""Count existing output frames for resume support.
Manifest-aware: reads .corridorkey
gitextract_m8ktm8q9/ ├── .dockerignore ├── .git-blame-ignore-revs ├── .github/ │ ├── ISSUE_TEMPLATE/ │ │ └── bug_report.yml │ ├── pull_request_template.md │ └── workflows/ │ ├── ci.yml │ └── docs.yml ├── .gitignore ├── .python-version ├── BiRefNetModule/ │ ├── checkpoints/ │ │ └── .gitkeep │ └── wrapper.py ├── CONTRIBUTING.md ├── ClipsForInference/ │ └── .gitkeep ├── CorridorKeyModule/ │ ├── IgnoredCheckpoints/ │ │ └── .gitkeep │ ├── README.md │ ├── __init__.py │ ├── backend.py │ ├── checkpoints/ │ │ └── .gitkeep │ ├── core/ │ │ ├── __init__.py │ │ ├── color_utils.py │ │ └── model_transformer.py │ └── inference_engine.py ├── CorridorKey_DRAG_CLIPS_HERE_local.bat ├── CorridorKey_DRAG_CLIPS_HERE_local.sh ├── Dockerfile ├── IgnoredClips/ │ └── .gitkeep ├── Install_CorridorKey_Linux_Mac.sh ├── Install_CorridorKey_Windows.bat ├── Install_GVM_Linux_Mac.sh ├── Install_GVM_Windows.bat ├── Install_VideoMaMa_Linux_Mac.sh ├── Install_VideoMaMa_Windows.bat ├── LICENSE ├── Output/ │ └── .gitkeep ├── README.md ├── RunGVMOnly.sh ├── RunInferenceOnly.sh ├── VideoMaMaInferenceModule/ │ ├── LICENSE.md │ ├── README.md │ ├── __init__.py │ ├── checkpoints/ │ │ └── .gitkeep │ ├── inference.py │ └── pipeline.py ├── backend/ │ ├── __init__.py │ ├── clip_state.py │ ├── errors.py │ ├── ffmpeg_tools.py │ ├── frame_io.py │ ├── job_queue.py │ ├── natural_sort.py │ ├── project.py │ ├── service.py │ └── validators.py ├── clip_manager.py ├── corridorkey_cli.py ├── device_utils.py ├── docker-compose.yml ├── docs/ │ ├── LLM_HANDOVER.md │ └── index.md ├── gvm_core/ │ ├── LICENSE.md │ ├── README.md │ ├── __init__.py │ ├── gvm/ │ │ ├── __init__.py │ │ ├── models/ │ │ │ ├── __init__.py │ │ │ └── unet_spatio_temporal_condition.py │ │ ├── pipelines/ │ │ │ └── pipeline_gvm.py │ │ └── utils/ │ │ ├── __init__.py │ │ └── inference_utils.py │ ├── weights/ │ │ └── .gitkeep │ └── wrapper.py ├── pyproject.toml ├── renovate.json ├── test_vram.py ├── tests/ │ ├── __init__.py │ ├── conftest.py │ ├── test_backend.py │ ├── test_cli.py │ ├── test_clip_manager.py │ ├── test_color_utils.py │ ├── test_device_utils.py │ ├── test_e2e_workflow.py │ ├── test_exr_gamma_bug_condition.py │ ├── test_exr_gamma_preservation.py │ ├── test_frame_io.py │ ├── test_gamma_consistency.py │ ├── test_imports.py │ ├── test_inference_engine.py │ ├── test_mlx_smoke.py │ ├── test_pbt_auto_download.py │ ├── test_pbt_backend_resolution.py │ ├── test_pbt_dep_preservation.py │ └── test_pyproject_structure.py └── zensical.toml
SYMBOL INDEX (655 symbols across 42 files)
FILE: BiRefNetModule/wrapper.py
class ImagePreprocessor (line 17) | class ImagePreprocessor:
method __init__ (line 18) | def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
method proc (line 27) | def proc(self, image: Image.Image) -> torch.Tensor:
class BiRefNetHandler (line 56) | class BiRefNetHandler:
method __init__ (line 57) | def __init__(self, device="cpu", usage="General"):
method cleanup (line 90) | def cleanup(self):
method process (line 106) | def process(self, input_path, alpha_output_dir=None, dilate_radius=0, ...
FILE: CorridorKeyModule/backend.py
function resolve_backend (line 32) | def resolve_backend(requested: str | None = None) -> str:
function _auto_detect_backend (line 62) | def _auto_detect_backend() -> str:
function _validate_mlx_available (line 106) | def _validate_mlx_available() -> None:
function _ensure_torch_checkpoint (line 120) | def _ensure_torch_checkpoint() -> Path:
function _discover_checkpoint (line 163) | def _discover_checkpoint(ext: str) -> Path:
function _wrap_mlx_output (line 189) | def _wrap_mlx_output(raw: dict, despill_strength: float, auto_despeckle:...
class _MLXEngineAdapter (line 238) | class _MLXEngineAdapter:
method __init__ (line 241) | def __init__(self, raw_engine):
method process_frame (line 245) | def process_frame(
function create_engine (line 290) | def create_engine(
FILE: CorridorKeyModule/core/color_utils.py
function _is_tensor (line 11) | def _is_tensor(x: np.ndarray | torch.Tensor) -> bool:
function _if_tensor (line 15) | def _if_tensor(is_tensor: bool, tensor_func: Callable, numpy_func: Calla...
function _power (line 19) | def _power(x: np.ndarray | torch.Tensor, exponent: float) -> np.ndarray ...
function _where (line 27) | def _where(
function _clamp (line 37) | def _clamp(x: np.ndarray | torch.Tensor, min: float) -> np.ndarray | tor...
function linear_to_srgb (line 50) | def linear_to_srgb(x: np.ndarray | torch.Tensor) -> np.ndarray | torch.T...
function srgb_to_linear (line 60) | def srgb_to_linear(x: np.ndarray | torch.Tensor) -> np.ndarray | torch.T...
function premultiply (line 70) | def premultiply(fg: np.ndarray | torch.Tensor, alpha: np.ndarray | torch...
function unpremultiply (line 79) | def unpremultiply(
function composite_straight (line 89) | def composite_straight(
function composite_premul (line 99) | def composite_premul(
function rgb_to_yuv (line 109) | def rgb_to_yuv(image: torch.Tensor) -> torch.Tensor:
function dilate_mask (line 146) | def dilate_mask(mask: np.ndarray | torch.Tensor, radius: int) -> np.ndar...
function apply_garbage_matte (line 181) | def apply_garbage_matte(
function despill (line 205) | def despill(
function clean_matte (line 250) | def clean_matte(alpha_np: np.ndarray, area_threshold: int = 300, dilatio...
function create_checkerboard (line 298) | def create_checkerboard(
FILE: CorridorKeyModule/core/model_transformer.py
class MLP (line 13) | class MLP(nn.Module):
method __init__ (line 16) | def __init__(self, input_dim: int = 2048, embed_dim: int = 768) -> None:
method forward (line 20) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class DecoderHead (line 24) | class DecoderHead(nn.Module):
method __init__ (line 25) | def __init__(
method forward (line 47) | def forward(self, features: list[torch.Tensor]) -> torch.Tensor:
class RefinerBlock (line 74) | class RefinerBlock(nn.Module):
method __init__ (line 79) | def __init__(self, channels: int, dilation: int = 1) -> None:
method forward (line 87) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class CNNRefinerModule (line 99) | class CNNRefinerModule(nn.Module):
method __init__ (line 106) | def __init__(self, in_channels: int = 7, hidden_channels: int = 64, ou...
method forward (line 129) | def forward(self, img: torch.Tensor, coarse_pred: torch.Tensor) -> tor...
class GreenFormer (line 145) | class GreenFormer(nn.Module):
method __init__ (line 146) | def __init__(
method _patch_input_layer (line 198) | def _patch_input_layer(self, in_channels: int) -> None:
method forward (line 242) | def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
FILE: CorridorKeyModule/inference_engine.py
class CorridorKeyEngine (line 19) | class CorridorKeyEngine:
method __init__ (line 20) | def __init__(
method _load_model (line 67) | def _load_model(self) -> GreenFormer:
method process_frame (line 125) | def process_frame(
FILE: VideoMaMaInferenceModule/inference.py
function load_videomama_model (line 26) | def load_videomama_model(base_model_path: Optional[str] = None, unet_che...
function extract_frames_from_video (line 66) | def extract_frames_from_video(video_path: str, max_frames: Optional[int]...
function run_inference (line 101) | def run_inference(
function save_video (line 178) | def save_video(frames: List[np.ndarray], output_path: str, fps: float):
FILE: VideoMaMaInferenceModule/pipeline.py
class StableVideoDiffusionPipelineOutput (line 62) | class StableVideoDiffusionPipelineOutput(BaseOutput):
class StableVideoDiffusionPipelineWithMask (line 73) | class StableVideoDiffusionPipelineWithMask(DiffusionPipeline):
method __init__ (line 83) | def __init__(
method _encode_image (line 103) | def _encode_image(
method _encode_vae_image (line 137) | def _encode_vae_image(
method _get_add_time_ids (line 148) | def _get_add_time_ids(
method decode_latents (line 168) | def decode_latents(self, latents: torch.Tensor, num_frames: int, decod...
method check_inputs (line 181) | def check_inputs(self, image, height, width):
method prepare_latents (line 191) | def prepare_latents(
method _encode_video_vae (line 230) | def _encode_video_vae(
method __call__ (line 248) | def __call__(
class StableVideoDiffusionPipelineOnestepWithMask (line 394) | class StableVideoDiffusionPipelineOnestepWithMask(DiffusionPipeline):
method __init__ (line 404) | def __init__(
method _encode_image (line 424) | def _encode_image(
method _encode_vae_image (line 458) | def _encode_vae_image(
method _get_add_time_ids (line 469) | def _get_add_time_ids(
method decode_latents (line 489) | def decode_latents(self, latents: torch.Tensor, num_frames: int, decod...
method check_inputs (line 502) | def check_inputs(self, image, height, width):
method prepare_latents (line 512) | def prepare_latents(
method _encode_video_vae (line 543) | def _encode_video_vae(
method __call__ (line 561) | def __call__(
class StableVideoDiffusionPipelineWithCrossAtnnMask (line 685) | class StableVideoDiffusionPipelineWithCrossAtnnMask(DiffusionPipeline):
method __init__ (line 689) | def __init__(
method _encode_image_vae (line 711) | def _encode_image_vae(self, image: torch.Tensor, device: Union[str, to...
method decode_latents (line 716) | def decode_latents(self, latents: torch.Tensor, num_frames: int, decod...
method _encode_video_vae (line 729) | def _encode_video_vae(
method __call__ (line 747) | def __call__(
class VideoInferencePipeline (line 838) | class VideoInferencePipeline:
method __init__ (line 846) | def __init__(self, base_model_path: str, unet_checkpoint_path: str, de...
method run (line 881) | def run(self, cond_frames, mask_frames, seed=42, mask_cond_mode="vae",...
method _pil_to_tensor (line 978) | def _pil_to_tensor(self, frames: list[Image.Image]):
method _tensor_to_vae_latent (line 983) | def _tensor_to_vae_latent(self, t: torch.Tensor):
method _get_add_time_ids (line 1001) | def _get_add_time_ids(self, fps, motion_bucket_id, noise_aug_strength,...
method _resize_with_antialiasing (line 1012) | def _resize_with_antialiasing(self, input_tensor, size, interpolation=...
FILE: backend/clip_state.py
class ClipState (line 44) | class ClipState(Enum):
class ClipAsset (line 65) | class ClipAsset:
method __post_init__ (line 72) | def __post_init__(self):
method _calculate_length (line 75) | def _calculate_length(self):
method get_frame_files (line 96) | def get_frame_files(self) -> list[str]:
class InOutRange (line 107) | class InOutRange:
method frame_count (line 114) | def frame_count(self) -> int:
method contains (line 117) | def contains(self, index: int) -> bool:
method to_dict (line 120) | def to_dict(self) -> dict:
method from_dict (line 124) | def from_dict(cls, d: dict) -> InOutRange:
class ClipEntry (line 129) | class ClipEntry:
method is_processing (line 146) | def is_processing(self) -> bool:
method set_processing (line 150) | def set_processing(self, value: bool) -> None:
method transition_to (line 154) | def transition_to(self, new_state: ClipState) -> None:
method set_error (line 164) | def set_error(self, message: str) -> None:
method output_dir (line 174) | def output_dir(self) -> str:
method has_outputs (line 178) | def has_outputs(self) -> bool:
method completed_frame_count (line 189) | def completed_frame_count(self) -> int:
method completed_stems (line 197) | def completed_stems(self) -> set[str]:
method _read_manifest (line 235) | def _read_manifest(self) -> dict | None:
method _resolve_original_path (line 249) | def _resolve_original_path(self) -> str | None:
method find_assets (line 262) | def find_assets(self) -> None:
method _resolve_state (line 332) | def _resolve_state(self) -> None:
function scan_project_clips (line 374) | def scan_project_clips(project_dir: str) -> list[ClipEntry]:
function scan_clips_dir (line 416) | def scan_clips_dir(
FILE: backend/errors.py
class CorridorKeyError (line 6) | class CorridorKeyError(Exception):
class ClipScanError (line 12) | class ClipScanError(CorridorKeyError):
class FrameMismatchError (line 18) | class FrameMismatchError(CorridorKeyError):
method __init__ (line 21) | def __init__(self, clip_name: str, input_count: int, alpha_count: int):
class FrameReadError (line 28) | class FrameReadError(CorridorKeyError):
method __init__ (line 31) | def __init__(self, clip_name: str, frame_index: int, path: str):
class WriteFailureError (line 38) | class WriteFailureError(CorridorKeyError):
method __init__ (line 41) | def __init__(self, clip_name: str, frame_index: int, path: str):
class MaskChannelError (line 48) | class MaskChannelError(CorridorKeyError):
method __init__ (line 51) | def __init__(self, clip_name: str, frame_index: int, channels: int):
class VRAMInsufficientError (line 58) | class VRAMInsufficientError(CorridorKeyError):
method __init__ (line 61) | def __init__(self, required_gb: float, available_gb: float):
class InvalidStateTransitionError (line 67) | class InvalidStateTransitionError(CorridorKeyError):
method __init__ (line 70) | def __init__(self, clip_name: str, current_state: str, target_state: s...
class JobCancelledError (line 77) | class JobCancelledError(CorridorKeyError):
method __init__ (line 80) | def __init__(self, clip_name: str, frame_index: int | None = None):
class FFmpegNotFoundError (line 89) | class FFmpegNotFoundError(CorridorKeyError):
method __init__ (line 92) | def __init__(self):
class ExtractionError (line 102) | class ExtractionError(CorridorKeyError):
method __init__ (line 105) | def __init__(self, clip_name: str, detail: str):
FILE: backend/ffmpeg_tools.py
function find_ffmpeg (line 34) | def find_ffmpeg() -> str | None:
function find_ffprobe (line 46) | def find_ffprobe() -> str | None:
function probe_video (line 58) | def probe_video(path: str) -> dict:
function extract_frames (line 133) | def extract_frames(
function stitch_video (line 291) | def stitch_video(
function write_video_metadata (line 383) | def write_video_metadata(clip_root: str, metadata: dict) -> None:
function read_video_metadata (line 395) | def read_video_metadata(clip_root: str) -> dict | None:
FILE: backend/frame_io.py
function read_image_frame (line 35) | def read_image_frame(fpath: str, gamma_correct_exr: bool = False) -> np....
function read_video_frame_at (line 70) | def read_video_frame_at(
function read_video_frames (line 98) | def read_video_frames(
function read_mask_frame (line 131) | def read_mask_frame(fpath: str, clip_name: str = "", frame_index: int = ...
function read_video_mask_at (line 155) | def read_video_mask_at(
FILE: backend/job_queue.py
class JobType (line 33) | class JobType(Enum):
class JobStatus (line 42) | class JobStatus(Enum):
class GPUJob (line 51) | class GPUJob:
method request_cancel (line 66) | def request_cancel(self) -> None:
method is_cancelled (line 71) | def is_cancelled(self) -> bool:
method check_cancelled (line 74) | def check_cancelled(self) -> None:
class GPUJobQueue (line 87) | class GPUJobQueue:
method __init__ (line 111) | def __init__(self):
method submit (line 123) | def submit(self, job: GPUJob) -> bool:
method next_job (line 163) | def next_job(self) -> GPUJob | None:
method start_job (line 170) | def start_job(self, job: GPUJob) -> None:
method complete_job (line 179) | def complete_job(self, job: GPUJob) -> None:
method fail_job (line 191) | def fail_job(self, job: GPUJob, error: str) -> None:
method mark_cancelled (line 204) | def mark_cancelled(self, job: GPUJob) -> None:
method cancel_job (line 218) | def cancel_job(self, job: GPUJob) -> None:
method cancel_current (line 232) | def cancel_current(self) -> None:
method cancel_all (line 238) | def cancel_all(self) -> None:
method report_progress (line 251) | def report_progress(self, clip_name: str, current: int, total: int) ->...
method report_warning (line 259) | def report_warning(self, message: str) -> None:
method find_job_by_id (line 265) | def find_job_by_id(self, job_id: str) -> GPUJob | None:
method clear_history (line 278) | def clear_history(self) -> None:
method remove_job (line 283) | def remove_job(self, job_id: str) -> None:
method has_pending (line 289) | def has_pending(self) -> bool:
method current_job (line 294) | def current_job(self) -> GPUJob | None:
method pending_count (line 299) | def pending_count(self) -> int:
method queue_snapshot (line 304) | def queue_snapshot(self) -> list[GPUJob]:
method history_snapshot (line 310) | def history_snapshot(self) -> list[GPUJob]:
method all_jobs_snapshot (line 316) | def all_jobs_snapshot(self) -> list[GPUJob]:
FILE: backend/natural_sort.py
function natural_sort_key (line 16) | def natural_sort_key(text: str) -> list[str | int]:
function natsorted (line 31) | def natsorted(items: list[str]) -> list[str]:
FILE: backend/project.py
function _dedupe_path (line 40) | def _dedupe_path(parent_dir: str, stem: str) -> tuple[str, str]:
function set_app_dir (line 62) | def set_app_dir(path: str) -> None:
function projects_root (line 68) | def projects_root() -> str:
function sanitize_stem (line 85) | def sanitize_stem(filename: str, max_len: int = 60) -> str:
function create_project (line 97) | def create_project(
function add_clips_to_project (line 175) | def add_clips_to_project(
function _create_clip_folder (line 214) | def _create_clip_folder(
function get_clip_dirs (line 256) | def get_clip_dirs(project_dir: str) -> list[str]:
function is_v2_project (line 273) | def is_v2_project(project_dir: str) -> bool:
function write_project_json (line 278) | def write_project_json(project_root: str, data: dict) -> None:
function read_project_json (line 287) | def read_project_json(project_root: str) -> dict | None:
function write_clip_json (line 300) | def write_clip_json(clip_root: str, data: dict) -> None:
function read_clip_json (line 309) | def read_clip_json(clip_root: str) -> dict | None:
function _read_clip_or_project_json (line 322) | def _read_clip_or_project_json(root: str) -> dict | None:
function get_display_name (line 330) | def get_display_name(root: str) -> str:
function set_display_name (line 341) | def set_display_name(root: str, name: str) -> None:
function save_in_out_range (line 353) | def save_in_out_range(clip_root: str, in_out) -> None:
function load_in_out_range (line 374) | def load_in_out_range(clip_root: str):
function is_video_file (line 387) | def is_video_file(filename: str) -> bool:
function is_image_file (line 392) | def is_image_file(filename: str) -> bool:
FILE: backend/service.py
class _ActiveModel (line 69) | class _ActiveModel(Enum):
class InferenceParams (line 79) | class InferenceParams:
method to_dict (line 88) | def to_dict(self) -> dict:
method from_dict (line 92) | def from_dict(cls, d: dict) -> "InferenceParams":
class OutputConfig (line 98) | class OutputConfig:
method to_dict (line 110) | def to_dict(self) -> dict:
method from_dict (line 114) | def from_dict(cls, d: dict) -> "OutputConfig":
method enabled_outputs (line 119) | def enabled_outputs(self) -> list[str]:
class FrameResult (line 134) | class FrameResult:
class CorridorKeyService (line 143) | class CorridorKeyService:
method __init__ (line 156) | def __init__(self):
method job_queue (line 167) | def job_queue(self) -> GPUJobQueue:
method detect_device (line 175) | def detect_device(self) -> str:
method get_vram_info (line 187) | def get_vram_info(self) -> dict[str, float]:
method _vram_allocated_mb (line 209) | def _vram_allocated_mb() -> float:
method _safe_offload (line 221) | def _safe_offload(obj: object) -> None:
method _ensure_model (line 240) | def _ensure_model(self, needed: _ActiveModel) -> None:
method _get_engine (line 284) | def _get_engine(self):
method _get_gvm (line 308) | def _get_gvm(self):
method _get_videomama_pipeline (line 323) | def _get_videomama_pipeline(self):
method unload_engines (line 339) | def unload_engines(self) -> None:
method scan_clips (line 358) | def scan_clips(
method get_clips_by_state (line 366) | def get_clips_by_state(
method _read_input_frame (line 376) | def _read_input_frame(
method _read_alpha_frame (line 410) | def _read_alpha_frame(
method _write_image (line 429) | def _write_image(
method _write_manifest (line 451) | def _write_manifest(
method _write_outputs (line 483) | def _write_outputs(
method run_inference (line 532) | def run_inference(
method is_engine_loaded (line 722) | def is_engine_loaded(self) -> bool:
method reprocess_single_frame (line 726) | def reprocess_single_frame(
method run_gvm (line 797) | def run_gvm(
method run_videomama (line 877) | def run_videomama(
method _load_frames_for_videomama (line 1049) | def _load_frames_for_videomama(
method _load_mask_frames_for_videomama (line 1080) | def _load_mask_frames_for_videomama(self, asset: ClipAsset, clip_name:...
FILE: backend/validators.py
function validate_frame_counts (line 23) | def validate_frame_counts(
function normalize_mask_channels (line 54) | def normalize_mask_channels(
function normalize_mask_dtype (line 82) | def normalize_mask_dtype(mask: np.ndarray) -> np.ndarray:
function validate_frame_read (line 96) | def validate_frame_read(
function validate_write (line 121) | def validate_write(
function ensure_output_dirs (line 142) | def ensure_output_dirs(clip_root: str) -> dict[str, str]:
FILE: clip_manager.py
class InferenceSettings (line 30) | class InferenceSettings:
function is_image_file (line 55) | def is_image_file(filename: str) -> bool:
function is_video_file (line 59) | def is_video_file(filename: str) -> bool:
function map_path (line 63) | def map_path(win_path: str) -> str:
class ClipAsset (line 83) | class ClipAsset:
method __init__ (line 84) | def __init__(self, path: str, asset_type: str) -> None:
method _calculate_length (line 90) | def _calculate_length(self) -> None:
class ClipEntry (line 103) | class ClipEntry:
method __init__ (line 104) | def __init__(self, name: str, root_path: str) -> None:
method find_assets (line 110) | def find_assets(self) -> None:
method validate_pair (line 180) | def validate_pair(self) -> None:
function get_gvm_processor (line 192) | def get_gvm_processor(device: str = "cpu") -> GVMProcessor:
function generate_alphas (line 205) | def generate_alphas(
function get_birefnet_usage_options (line 290) | def get_birefnet_usage_options():
function run_birefnet (line 294) | def run_birefnet(
function run_videomama (line 352) | def run_videomama(
function run_inference (line 596) | def run_inference(
function organize_target (line 852) | def organize_target(target_dir: str) -> None:
function organize_clips (line 912) | def organize_clips(clips_dir: str) -> None:
function scan_clips (line 955) | def scan_clips() -> list[ClipEntry]:
FILE: corridorkey_cli.py
function _configure_environment (line 64) | def _configure_environment() -> None:
class ProgressContext (line 81) | class ProgressContext:
method __init__ (line 89) | def __init__(self) -> None:
method __enter__ (line 100) | def __enter__(self) -> "ProgressContext":
method __exit__ (line 104) | def __exit__(self, *exc: object) -> None:
method on_clip_start (line 107) | def on_clip_start(self, clip_name: str, num_frames: int) -> None:
method on_frame_complete (line 113) | def on_frame_complete(self, frame_idx: int, num_frames: int) -> None:
function _on_clip_start_log_only (line 119) | def _on_clip_start_log_only(clip_name: str, total_clips: int) -> None:
function _prompt_inference_settings (line 133) | def _prompt_inference_settings(
function app_callback (line 207) | def app_callback(
function list_clips_cmd (line 227) | def list_clips_cmd(ctx: typer.Context) -> None:
function generate_alphas_cmd (line 233) | def generate_alphas_cmd(ctx: typer.Context) -> None:
function run_inference_cmd (line 242) | def run_inference_cmd(
function wizard (line 321) | def wizard(
function interactive_wizard (line 334) | def interactive_wizard(win_path: str, device: str | None = None) -> None:
function main (line 574) | def main() -> None:
FILE: device_utils.py
function detect_best_device (line 14) | def detect_best_device() -> str:
function resolve_device (line 26) | def resolve_device(requested: str | None = None) -> str:
function clear_device_cache (line 70) | def clear_device_cache(device: torch.device | str) -> None:
FILE: gvm_core/gvm/models/unet_spatio_temporal_condition.py
class UNetSpatioTemporalConditionModel (line 25) | class UNetSpatioTemporalConditionModel(
method __init__ (line 69) | def __init__(
method _set_class_embedding (line 274) | def _set_class_embedding(
method get_class_embed (line 311) | def get_class_embed(self, sample: torch.Tensor, class_labels: Optional...
method attn_processors (line 329) | def attn_processors(self) -> Dict[str, AttentionProcessor]:
method set_attn_processor (line 358) | def set_attn_processor(
method set_default_attn_processor (line 394) | def set_default_attn_processor(self):
method _set_gradient_checkpointing (line 410) | def _set_gradient_checkpointing(self, module, value=False):
method enable_forward_chunking (line 415) | def enable_forward_chunking(
method forward (line 448) | def forward(
FILE: gvm_core/gvm/pipelines/pipeline_gvm.py
class GVMLoraLoader (line 25) | class GVMLoraLoader(StableDiffusionLoraLoaderMixin):
method __init__ (line 27) | def __init__(self, *args, **kwargs):
method load_lora_weights (line 30) | def load_lora_weights(
class GVMOutput (line 45) | class GVMOutput(BaseOutput):
class GVMPipeline (line 57) | class GVMPipeline(DiffusionPipeline, GVMLoraLoader):
method __init__ (line 58) | def __init__(self, vae, unet, scheduler):
method encode (line 64) | def encode(self, input):
method decode (line 72) | def decode(self, latents, decode_chunk_size=16):
method single_infer (line 94) | def single_infer(self, rgb, position_ids=None, num_inference_steps=Non...
method __call__ (line 136) | def __call__(
FILE: gvm_core/gvm/utils/inference_utils.py
class VideoReader (line 12) | class VideoReader(Dataset):
method __init__ (line 13) | def __init__(self, path, max_frames=None, transform=None):
method frame_rate (line 20) | def frame_rate(self):
method origin_shape (line 24) | def origin_shape(self):
method __len__ (line 27) | def __len__(self):
method __getitem__ (line 33) | def __getitem__(self, idx):
class VideoWriter (line 41) | class VideoWriter:
method __init__ (line 42) | def __init__(self, path, frame_rate, bit_rate=1000000):
method write (line 49) | def write(self, frames):
method write_numpy (line 63) | def write_numpy(self, frames):
method close (line 74) | def close(self):
class ImageSequenceReader (line 79) | class ImageSequenceReader(Dataset):
method __init__ (line 80) | def __init__(self, path, transform=None):
method origin_shape (line 86) | def origin_shape(self):
method __len__ (line 92) | def __len__(self):
method __getitem__ (line 95) | def __getitem__(self, idx):
class ImageSequenceWriter (line 143) | class ImageSequenceWriter:
method __init__ (line 144) | def __init__(self, path, extension='jpg'):
method write (line 150) | def write(self, frames, filenames=None):
method close (line 162) | def close(self):
FILE: gvm_core/wrapper.py
function seed_all (line 26) | def seed_all(seed: int = 0):
function impad_multi (line 38) | def impad_multi(img, multiple=32):
function sequence_collate_fn (line 55) | def sequence_collate_fn(examples):
class GVMProcessor (line 61) | class GVMProcessor:
method __init__ (line 62) | def __init__(self,
method process_sequence (line 106) | def process_sequence(self, input_path, output_dir,
FILE: test_vram.py
function process_frame (line 9) | def process_frame(engine):
function test_vram (line 16) | def test_vram():
FILE: tests/conftest.py
function _has_gpu (line 13) | def _has_gpu():
function _has_mlx (line 23) | def _has_mlx():
function pytest_collection_modifyitems (line 35) | def pytest_collection_modifyitems(config, items):
function sample_frame_rgb (line 56) | def sample_frame_rgb():
function sample_mask (line 63) | def sample_mask():
function tmp_clip_dir (line 77) | def tmp_clip_dir(tmp_path):
function mock_greenformer (line 139) | def mock_greenformer():
function silent_backend_injection (line 166) | def silent_backend_injection(monkeypatch):
function stage_shot (line 189) | def stage_shot(tmp_path):
function sandbox_clip_manager (line 224) | def sandbox_clip_manager(tmp_path, monkeypatch):
FILE: tests/test_backend.py
class TestResolveBackend (line 26) | class TestResolveBackend:
method test_explicit_torch (line 27) | def test_explicit_torch(self):
method test_explicit_mlx_on_non_apple_raises (line 30) | def test_explicit_mlx_on_non_apple_raises(self):
method test_env_var_torch (line 36) | def test_env_var_torch(self):
method test_auto_non_darwin (line 41) | def test_auto_non_darwin(self):
method test_auto_darwin_no_mlx_package (line 46) | def test_auto_darwin_no_mlx_package(self):
method test_unknown_backend_raises (line 67) | def test_unknown_backend_raises(self):
class TestDiscoverCheckpoint (line 75) | class TestDiscoverCheckpoint:
method test_exactly_one (line 76) | def test_exactly_one(self, tmp_path):
method test_zero_torch_triggers_auto_download (line 83) | def test_zero_torch_triggers_auto_download(self, tmp_path):
method test_zero_torch_download_failure_raises_runtime_error (line 98) | def test_zero_torch_download_failure_raises_runtime_error(self, tmp_pa...
method test_zero_safetensors_with_cross_reference (line 108) | def test_zero_safetensors_with_cross_reference(self, tmp_path):
method test_multiple_raises (line 115) | def test_multiple_raises(self, tmp_path):
method test_safetensors (line 122) | def test_safetensors(self, tmp_path):
method test_ensure_torch_checkpoint_happy_path (line 129) | def test_ensure_torch_checkpoint_happy_path(self, tmp_path):
method test_skip_when_present (line 147) | def test_skip_when_present(self, tmp_path):
method test_mlx_not_triggered (line 157) | def test_mlx_not_triggered(self, tmp_path):
method test_network_error_wrapping (line 165) | def test_network_error_wrapping(self, tmp_path):
method test_disk_space_error (line 176) | def test_disk_space_error(self, tmp_path):
method test_logging_on_download (line 191) | def test_logging_on_download(self, tmp_path, caplog):
class TestWrapMlxOutput (line 209) | class TestWrapMlxOutput:
method mlx_raw_output (line 211) | def mlx_raw_output(self):
method test_output_keys (line 222) | def test_output_keys(self, mlx_raw_output):
method test_alpha_shape_dtype (line 226) | def test_alpha_shape_dtype(self, mlx_raw_output):
method test_fg_shape_dtype (line 233) | def test_fg_shape_dtype(self, mlx_raw_output):
method test_processed_shape_dtype (line 238) | def test_processed_shape_dtype(self, mlx_raw_output):
method test_comp_shape_dtype (line 243) | def test_comp_shape_dtype(self, mlx_raw_output):
method test_value_ranges (line 248) | def test_value_ranges(self, mlx_raw_output):
FILE: tests/test_cli.py
class TestHelpOutput (line 23) | class TestHelpOutput:
method test_main_help (line 24) | def test_main_help(self):
method test_list_clips_help (line 32) | def test_list_clips_help(self):
method test_generate_alphas_help (line 36) | def test_generate_alphas_help(self):
method test_run_inference_help (line 40) | def test_run_inference_help(self):
method test_wizard_help (line 44) | def test_wizard_help(self):
class TestInvalidArgs (line 54) | class TestInvalidArgs:
method test_wizard_requires_path (line 55) | def test_wizard_requires_path(self):
method test_unknown_subcommand (line 59) | def test_unknown_subcommand(self):
class TestInferenceSettings (line 69) | class TestInferenceSettings:
method test_defaults (line 70) | def test_defaults(self):
method test_custom_values (line 78) | def test_custom_values(self):
class TestCallbackProtocol (line 98) | class TestCallbackProtocol:
method test_run_inference_passes_callbacks (line 102) | def test_run_inference_passes_callbacks(self, mock_prompt, mock_run, m...
method test_callback_signatures (line 117) | def test_callback_signatures(self):
class TestListClips (line 137) | class TestListClips:
method test_list_clips_calls_scan (line 139) | def test_list_clips_calls_scan(self, mock_scan):
class TestNonInteractiveFlags (line 151) | class TestNonInteractiveFlags:
method test_all_flags_skips_prompts (line 154) | def test_all_flags_skips_prompts(self, mock_run, mock_scan):
method test_srgb_flag (line 185) | def test_srgb_flag(self, mock_run, mock_scan):
method test_despill_clamped_to_range (line 211) | def test_despill_clamped_to_range(self, mock_run, mock_scan):
method test_run_inference_help_shows_flags (line 234) | def test_run_inference_help_shows_flags(self):
method test_skip_existing_passed_through (line 247) | def test_skip_existing_passed_through(self, mock_run, mock_scan):
FILE: tests/test_clip_manager.py
class TestFileTypeDetection (line 37) | class TestFileTypeDetection:
method test_image_extensions_recognized (line 56) | def test_image_extensions_recognized(self, filename):
method test_video_extensions_recognized (line 68) | def test_video_extensions_recognized(self, filename):
method test_non_media_rejected (line 81) | def test_non_media_rejected(self, filename):
method test_image_is_not_video (line 85) | def test_image_is_not_video(self):
method test_video_is_not_image (line 90) | def test_video_is_not_image(self):
class TestMapPath (line 100) | class TestMapPath:
method test_basic_mapping (line 107) | def test_basic_mapping(self):
method test_case_insensitive_drive_letter (line 111) | def test_case_insensitive_drive_letter(self):
method test_trailing_whitespace_stripped (line 115) | def test_trailing_whitespace_stripped(self):
method test_backslashes_converted (line 119) | def test_backslashes_converted(self):
method test_non_v_drive_passthrough (line 123) | def test_non_v_drive_passthrough(self):
method test_drive_root_only (line 128) | def test_drive_root_only(self):
class TestClipAsset (line 138) | class TestClipAsset:
method test_sequence_frame_count (line 141) | def test_sequence_frame_count(self, tmp_path):
method test_sequence_ignores_non_image_files (line 152) | def test_sequence_ignores_non_image_files(self, tmp_path):
method test_empty_sequence (line 164) | def test_empty_sequence(self, tmp_path):
class TestClipEntryFindAssets (line 177) | class TestClipEntryFindAssets:
method test_finds_image_sequence_input (line 184) | def test_finds_image_sequence_input(self, tmp_clip_dir):
method test_finds_alpha_hint (line 192) | def test_finds_alpha_hint(self, tmp_clip_dir):
method test_empty_alpha_hint_is_none (line 200) | def test_empty_alpha_hint_is_none(self, tmp_clip_dir):
method test_missing_input_raises (line 207) | def test_missing_input_raises(self, tmp_path):
method test_empty_input_dir_raises (line 215) | def test_empty_input_dir_raises(self, tmp_path):
method test_validate_pair_frame_count_mismatch (line 223) | def test_validate_pair_frame_count_mismatch(self, tmp_path):
method test_validate_pair_matching_counts_ok (line 243) | def test_validate_pair_matching_counts_ok(self, tmp_clip_dir):
class TestGenerateAlphas (line 255) | class TestGenerateAlphas:
method test_all_clips_valid_skips_generation (line 261) | def test_all_clips_valid_skips_generation(self, caplog):
method test_gvm_missing_exits_gracefully (line 275) | def test_gvm_missing_exits_gracefully(self, mock_get_processor, caplog):
method test_existing_alpha_dir_is_cleaned (line 291) | def test_existing_alpha_dir_is_cleaned(self, _mock_gvm, tmp_path):
method test_naming_remap_sequence (line 312) | def test_naming_remap_sequence(self, mock_get_processor, tmp_path):
method test_naming_remap_video (line 346) | def test_naming_remap_video(self, mock_get_processor, tmp_path):
method test_empty_output_logs_error (line 378) | def test_empty_output_logs_error(self, mock_get_processor, tmp_path, c...
class TestVideoMaMa (line 405) | class TestVideoMaMa:
method test_videomama_skips_if_sequence_exists (line 406) | def test_videomama_skips_if_sequence_exists(self, stage_shot, caplog):
method test_videomama_processes_valid_candidate (line 426) | def test_videomama_processes_valid_candidate(self, stage_shot):
method test_videomama_skips_if_input_missing (line 439) | def test_videomama_skips_if_input_missing(self, tmp_path):
method test_videomama_skips_if_mask_missing (line 450) | def test_videomama_skips_if_mask_missing(self, stage_shot, caplog):
method test_videomama_mask_thresholding (line 469) | def test_videomama_mask_thresholding(self, stage_shot):
method test_videomama_rgba_to_rgb_conversion (line 483) | def test_videomama_rgba_to_rgb_conversion(self, stage_shot):
method test_videomama_exr_gamma_handling (line 497) | def test_videomama_exr_gamma_handling(self, stage_shot):
method test_safety_removes_file_blocking_dir (line 512) | def test_safety_removes_file_blocking_dir(self, stage_shot):
method test_videomama_multiple_clips_batch (line 526) | def test_videomama_multiple_clips_batch(self, stage_shot):
method test_videomama_upgrades_video_alpha (line 541) | def test_videomama_upgrades_video_alpha(self, stage_shot):
method test_videomama_handles_invalid_image_load (line 552) | def test_videomama_handles_invalid_image_load(self, stage_shot, caplog):
method test_videomama_priority_folder_over_video (line 574) | def test_videomama_priority_folder_over_video(self, stage_shot):
method test_loop_chunking_logic (line 587) | def test_loop_chunking_logic(self, tmp_path):
method test_videomama_mask_from_video (line 611) | def test_videomama_mask_from_video(self, stage_shot):
method test_videomama_cleanup_on_failure (line 622) | def test_videomama_cleanup_on_failure(self, stage_shot, caplog):
class TestOrganizeTarget (line 657) | class TestOrganizeTarget:
method test_creates_hint_directories (line 663) | def test_creates_hint_directories(self, tmp_path):
method test_existing_hint_dirs_preserved (line 675) | def test_existing_hint_dirs_preserved(self, tmp_clip_dir):
method test_moves_loose_images_to_input (line 685) | def test_moves_loose_images_to_input(self, tmp_path):
class TestOrganizeClips (line 707) | class TestOrganizeClips:
method test_organize_loose_video_file (line 712) | def test_organize_loose_video_file(self, tmp_path):
method test_skips_video_if_folder_exists (line 735) | def test_skips_video_if_folder_exists(self, tmp_path, caplog):
method test_ignores_protected_folders (line 755) | def test_ignores_protected_folders(self, tmp_path):
method test_handles_nonexistent_directory (line 776) | def test_handles_nonexistent_directory(self, caplog):
method test_batch_organization_mix (line 790) | def test_batch_organization_mix(self, tmp_path):
class TestScanClips (line 821) | class TestScanClips:
method test_creates_clips_dir_and_returns_empty_if_missing (line 828) | def test_creates_clips_dir_and_returns_empty_if_missing(self, tmp_path...
method test_returns_clips_with_valid_input (line 840) | def test_returns_clips_with_valid_input(self, tmp_clip_dir, monkeypatch):
method test_excludes_frame_count_mismatch (line 851) | def test_excludes_frame_count_mismatch(self, tmp_clip_dir, monkeypatch):
method test_skips_hidden_and_underscore_dirs (line 871) | def test_skips_hidden_and_underscore_dirs(self, tmp_clip_dir, monkeypa...
method test_noise_filter_skips_hidden_folders (line 885) | def test_noise_filter_skips_hidden_folders(self, sandbox_clip_manager):
method test_scanner_handles_multiple_shots (line 903) | def test_scanner_handles_multiple_shots(self, sandbox_clip_manager):
method test_ideal_organization_loose_videos (line 920) | def test_ideal_organization_loose_videos(self, sandbox_clip_manager):
method test_organization_skips_existing_folders (line 934) | def test_organization_skips_existing_folders(self, sandbox_clip_manage...
method test_batch_processing_mix (line 947) | def test_batch_processing_mix(self, sandbox_clip_manager):
method test_nonexistent_directory_logging (line 960) | def test_nonexistent_directory_logging(self, caplog):
FILE: tests/test_color_utils.py
function _to_np (line 21) | def _to_np(x):
function _to_torch (line 26) | def _to_torch(x):
class TestSrgbLinearConversion (line 36) | class TestSrgbLinearConversion:
method test_identity_values_numpy (line 45) | def test_identity_values_numpy(self, value):
method test_identity_values_torch (line 51) | def test_identity_values_torch(self, value):
method test_mid_gray_numpy (line 57) | def test_mid_gray_numpy(self):
method test_mid_gray_torch (line 62) | def test_mid_gray_torch(self):
method test_roundtrip_numpy (line 69) | def test_roundtrip_numpy(self, value):
method test_roundtrip_torch (line 75) | def test_roundtrip_torch(self, value):
method test_breakpoint_continuity_linear_to_srgb (line 81) | def test_breakpoint_continuity_linear_to_srgb(self):
method test_breakpoint_continuity_srgb_to_linear (line 91) | def test_breakpoint_continuity_srgb_to_linear(self):
method test_negative_clamped_linear_to_srgb_numpy (line 100) | def test_negative_clamped_linear_to_srgb_numpy(self):
method test_negative_clamped_linear_to_srgb_torch (line 104) | def test_negative_clamped_linear_to_srgb_torch(self):
method test_negative_clamped_srgb_to_linear_numpy (line 108) | def test_negative_clamped_srgb_to_linear_numpy(self):
method test_negative_clamped_srgb_to_linear_torch (line 112) | def test_negative_clamped_srgb_to_linear_torch(self):
method test_vectorized_numpy (line 117) | def test_vectorized_numpy(self):
method test_vectorized_torch (line 124) | def test_vectorized_torch(self):
class TestPremultiply (line 137) | class TestPremultiply:
method test_roundtrip_numpy (line 143) | def test_roundtrip_numpy(self):
method test_roundtrip_torch (line 150) | def test_roundtrip_torch(self):
method test_output_bounded_by_fg_numpy (line 157) | def test_output_bounded_by_fg_numpy(self):
method test_output_bounded_by_fg_torch (line 164) | def test_output_bounded_by_fg_torch(self):
method test_zero_alpha_numpy (line 170) | def test_zero_alpha_numpy(self):
method test_one_alpha_numpy (line 177) | def test_one_alpha_numpy(self):
class TestCompositing (line 190) | class TestCompositing:
method test_straight_vs_premul_equivalence_numpy (line 197) | def test_straight_vs_premul_equivalence_numpy(self):
method test_straight_vs_premul_equivalence_torch (line 208) | def test_straight_vs_premul_equivalence_torch(self):
method test_alpha_zero_shows_background (line 219) | def test_alpha_zero_shows_background(self):
method test_alpha_one_shows_foreground (line 226) | def test_alpha_one_shows_foreground(self):
class TestDespill (line 239) | class TestDespill:
method test_pure_green_reduced_average_mode_numpy (line 246) | def test_pure_green_reduced_average_mode_numpy(self):
method test_pure_green_reduced_max_mode_numpy (line 253) | def test_pure_green_reduced_max_mode_numpy(self):
method test_pure_red_unchanged_numpy (line 259) | def test_pure_red_unchanged_numpy(self):
method test_strength_zero_is_noop_numpy (line 265) | def test_strength_zero_is_noop_numpy(self):
method test_partial_green_average_mode_numpy (line 271) | def test_partial_green_average_mode_numpy(self):
method test_max_mode_higher_limit_than_average (line 279) | def test_max_mode_higher_limit_than_average(self):
method test_fractional_strength_interpolates (line 287) | def test_fractional_strength_interpolates(self):
method test_despill_torch (line 299) | def test_despill_torch(self):
method test_green_below_limit_unchanged_numpy (line 307) | def test_green_below_limit_unchanged_numpy(self):
class TestCleanMatte (line 327) | class TestCleanMatte:
method test_large_blob_preserved (line 334) | def test_large_blob_preserved(self):
method test_small_blob_removed (line 342) | def test_small_blob_removed(self):
method test_mixed_blobs (line 349) | def test_mixed_blobs(self):
method test_3d_input_preserved (line 361) | def test_3d_input_preserved(self):
class TestCheckerboard (line 375) | class TestCheckerboard:
method test_output_shape (line 378) | def test_output_shape(self):
method test_output_range (line 382) | def test_output_range(self):
method test_uses_specified_colors (line 387) | def test_uses_specified_colors(self):
class TestRgbToYuv (line 398) | class TestRgbToYuv:
method test_pure_white_bchw (line 408) | def test_pure_white_bchw(self):
method test_pure_red_known_values (line 419) | def test_pure_red_known_values(self):
method test_chw_layout (line 431) | def test_chw_layout(self):
method test_last_dim_layout (line 440) | def test_last_dim_layout(self):
method test_rejects_numpy (line 452) | def test_rejects_numpy(self):
class TestDilateMask (line 464) | class TestDilateMask:
method test_radius_zero_noop_numpy (line 470) | def test_radius_zero_noop_numpy(self):
method test_radius_zero_noop_torch (line 476) | def test_radius_zero_noop_torch(self):
method test_dilation_expands_numpy (line 482) | def test_dilation_expands_numpy(self):
method test_dilation_expands_torch (line 492) | def test_dilation_expands_torch(self):
method test_preserves_2d_shape_numpy (line 501) | def test_preserves_2d_shape_numpy(self):
method test_preserves_2d_shape_torch (line 506) | def test_preserves_2d_shape_torch(self):
method test_preserves_3d_shape_torch (line 511) | def test_preserves_3d_shape_torch(self):
class TestApplyGarbageMatte (line 523) | class TestApplyGarbageMatte:
method test_none_input_passthrough (line 529) | def test_none_input_passthrough(self):
method test_zeros_outside_garbage_region (line 536) | def test_zeros_outside_garbage_region(self):
method test_3d_matte_with_2d_garbage (line 547) | def test_3d_matte_with_2d_garbage(self):
FILE: tests/test_device_utils.py
function _patch_gpu (line 25) | def _patch_gpu(monkeypatch, *, cuda=False, mps=False):
class TestDetectBestDevice (line 39) | class TestDetectBestDevice:
method test_returns_cuda_when_available (line 42) | def test_returns_cuda_when_available(self, monkeypatch):
method test_returns_mps_when_no_cuda (line 46) | def test_returns_mps_when_no_cuda(self, monkeypatch):
method test_returns_cpu_when_nothing (line 50) | def test_returns_cpu_when_nothing(self, monkeypatch):
class TestResolveDevice (line 60) | class TestResolveDevice:
method test_none_triggers_auto_detect (line 65) | def test_none_triggers_auto_detect(self, monkeypatch):
method test_auto_string_triggers_auto_detect (line 70) | def test_auto_string_triggers_auto_detect(self, monkeypatch):
method test_env_var_used_when_no_cli_arg (line 77) | def test_env_var_used_when_no_cli_arg(self, monkeypatch):
method test_env_var_auto_triggers_detect (line 82) | def test_env_var_auto_triggers_detect(self, monkeypatch):
method test_cli_arg_overrides_env_var (line 89) | def test_cli_arg_overrides_env_var(self, monkeypatch):
method test_explicit_cuda (line 96) | def test_explicit_cuda(self, monkeypatch):
method test_explicit_mps (line 100) | def test_explicit_mps(self, monkeypatch):
method test_explicit_cpu (line 104) | def test_explicit_cpu(self, monkeypatch):
method test_case_insensitive (line 107) | def test_case_insensitive(self, monkeypatch):
method test_cuda_unavailable_raises (line 112) | def test_cuda_unavailable_raises(self, monkeypatch):
method test_mps_no_backend_raises (line 117) | def test_mps_no_backend_raises(self, monkeypatch):
method test_mps_unavailable_raises (line 126) | def test_mps_unavailable_raises(self, monkeypatch):
method test_invalid_device_raises (line 133) | def test_invalid_device_raises(self, monkeypatch):
class TestClearDeviceCache (line 143) | class TestClearDeviceCache:
method test_cuda_clears_cache (line 146) | def test_cuda_clears_cache(self, monkeypatch):
method test_mps_clears_cache (line 152) | def test_mps_clears_cache(self, monkeypatch):
method test_cpu_is_noop (line 158) | def test_cpu_is_noop(self):
method test_accepts_torch_device_object (line 162) | def test_accepts_torch_device_object(self, monkeypatch):
method test_accepts_mps_device_object (line 168) | def test_accepts_mps_device_object(self, monkeypatch):
FILE: tests/test_e2e_workflow.py
function _fake_result (line 25) | def _fake_result(h: int = 4, w: int = 4) -> dict:
class TestE2EInferenceWorkflow (line 40) | class TestE2EInferenceWorkflow:
method test_output_directories_created (line 48) | def test_output_directories_created(self, tmp_clip_dir, monkeypatch):
method test_output_files_written_per_frame (line 70) | def test_output_files_written_per_frame(self, tmp_clip_dir, monkeypatch):
method test_clip_without_alpha_skipped (line 96) | def test_clip_without_alpha_skipped(self, tmp_clip_dir, monkeypatch):
FILE: tests/test_exr_gamma_bug_condition.py
function _write_exr (line 64) | def _write_exr(path: str, rgb_data: np.ndarray) -> None:
class TestDefect3NaivePowVsPiecewiseSRGB (line 75) | class TestDefect3NaivePowVsPiecewiseSRGB:
method test_gamma_corrected_exr_matches_piecewise_srgb (line 84) | def test_gamma_corrected_exr_matches_piecewise_srgb(self, data: np.nda...
method test_threshold_region_divergence (line 120) | def test_threshold_region_divergence(self, pixel_val: float) -> None:
class TestDefect1And2RunInferenceEXRPath (line 156) | class TestDefect1And2RunInferenceEXRPath:
method test_exr_srgb_frame_is_gamma_corrected (line 167) | def test_exr_srgb_frame_is_gamma_corrected(self, data: np.ndarray) -> ...
FILE: tests/test_exr_gamma_preservation.py
function _write_exr (line 61) | def _write_exr(path: str, rgb_data: np.ndarray) -> None:
function _write_png (line 67) | def _write_png(path: str, rgb_data: np.ndarray) -> None:
class TestPreservationLinearEXRRead (line 78) | class TestPreservationLinearEXRRead:
method test_exr_default_returns_raw_linear_data (line 88) | def test_exr_default_returns_raw_linear_data(self, data: np.ndarray) -...
method test_exr_explicit_false_returns_raw_linear_data (line 121) | def test_exr_explicit_false_returns_raw_linear_data(self, data: np.nda...
class TestPreservationStandardImageRead (line 152) | class TestPreservationStandardImageRead:
method test_png_read_returns_uint8_normalized (line 161) | def test_png_read_returns_uint8_normalized(self, data: np.ndarray) -> ...
method test_png_result_dtype_and_range (line 189) | def test_png_result_dtype_and_range(self, data: np.ndarray) -> None:
class TestPreservationLinearEXRInference (line 211) | class TestPreservationLinearEXRInference:
method test_linear_exr_passes_through_unchanged (line 221) | def test_linear_exr_passes_through_unchanged(self, data: np.ndarray) -...
FILE: tests/test_frame_io.py
class TestReadVideoFrameAtNegativeIndex (line 12) | class TestReadVideoFrameAtNegativeIndex:
method test_negative_one_returns_none (line 15) | def test_negative_one_returns_none(self, tmp_path):
method test_large_negative_returns_none (line 21) | def test_large_negative_returns_none(self, tmp_path):
method test_zero_does_not_trigger_guard (line 26) | def test_zero_does_not_trigger_guard(self, tmp_path):
class TestReadVideoMaskAtNegativeIndex (line 38) | class TestReadVideoMaskAtNegativeIndex:
method test_negative_one_returns_none (line 41) | def test_negative_one_returns_none(self, tmp_path):
method test_large_negative_returns_none (line 46) | def test_large_negative_returns_none(self, tmp_path):
method test_zero_does_not_trigger_guard (line 51) | def test_zero_does_not_trigger_guard(self, tmp_path):
FILE: tests/test_gamma_consistency.py
class TestGammaInconsistency (line 41) | class TestGammaInconsistency:
method test_linear_to_srgb_differs_from_gamma_22 (line 47) | def test_linear_to_srgb_differs_from_gamma_22(self):
method test_srgb_to_linear_differs_from_gamma_22 (line 64) | def test_srgb_to_linear_differs_from_gamma_22(self):
class TestGammaDivergenceMagnitude (line 82) | class TestGammaDivergenceMagnitude:
method test_divergence_at_known_values (line 104) | def test_divergence_at_known_values(self, linear_val, expected_min_diff):
method test_both_methods_agree_at_zero_and_one (line 117) | def test_both_methods_agree_at_zero_and_one(self):
method test_worst_case_divergence_in_darks (line 125) | def test_worst_case_divergence_in_darks(self):
FILE: tests/test_imports.py
function test_import_corridorkey_module (line 8) | def test_import_corridorkey_module():
function test_import_color_utils (line 12) | def test_import_color_utils():
function test_import_inference_engine (line 16) | def test_import_inference_engine():
function test_import_model_transformer (line 20) | def test_import_model_transformer():
function test_import_gvm_core (line 24) | def test_import_gvm_core():
function test_import_gvm_wrapper (line 28) | def test_import_gvm_wrapper():
function test_import_videomama (line 32) | def test_import_videomama():
function test_import_videomama_inference (line 36) | def test_import_videomama_inference():
FILE: tests/test_inference_engine.py
function _make_engine_with_mock (line 27) | def _make_engine_with_mock(mock_greenformer, img_size=64):
class TestProcessFrameOutputs (line 53) | class TestProcessFrameOutputs:
method test_output_keys (line 56) | def test_output_keys(self, sample_frame_rgb, sample_mask, mock_greenfo...
method test_output_shapes_match_input (line 66) | def test_output_shapes_match_input(self, sample_frame_rgb, sample_mask...
method test_output_dtype_float32 (line 77) | def test_output_dtype_float32(self, sample_frame_rgb, sample_mask, moc...
method test_alpha_output_range_is_zero_to_one (line 85) | def test_alpha_output_range_is_zero_to_one(self, sample_frame_rgb, sam...
method test_fg_output_range_is_zero_to_one (line 93) | def test_fg_output_range_is_zero_to_one(self, sample_frame_rgb, sample...
class TestProcessFrameColorSpace (line 107) | class TestProcessFrameColorSpace:
method test_srgb_input_default (line 115) | def test_srgb_input_default(self, sample_frame_rgb, sample_mask, mock_...
method test_linear_input_path (line 122) | def test_linear_input_path(self, sample_frame_rgb, sample_mask, mock_g...
method test_uint8_input_normalized (line 128) | def test_uint8_input_normalized(self, sample_mask, mock_greenformer):
method test_model_called_exactly_once (line 136) | def test_model_called_exactly_once(self, sample_frame_rgb, sample_mask...
class TestProcessFramePostProcessing (line 151) | class TestProcessFramePostProcessing:
method test_despill_strength_reduces_green_in_spill_pixels (line 154) | def test_despill_strength_reduces_green_in_spill_pixels(self, sample_f...
method test_auto_despeckle_toggle (line 197) | def test_auto_despeckle_toggle(self, sample_frame_rgb, sample_mask, mo...
method test_processed_is_linear_premul_rgba (line 203) | def test_processed_is_linear_premul_rgba(self, sample_frame_rgb, sampl...
method test_mask_2d_vs_3d_input (line 226) | def test_mask_2d_vs_3d_input(self, sample_frame_rgb, mock_greenformer):
method test_refiner_scale_parameter_accepted (line 238) | def test_refiner_scale_parameter_accepted(self, sample_frame_rgb, samp...
class TestNvidiaGPUProcess (line 250) | class TestNvidiaGPUProcess:
method test_process_frame_on_gpu (line 252) | def test_process_frame_on_gpu(self, sample_frame_rgb, sample_mask, moc...
FILE: tests/test_mlx_smoke.py
function mlx_engine (line 10) | def mlx_engine():
function test_mlx_smoke_2048 (line 17) | def test_mlx_smoke_2048(mlx_engine):
FILE: tests/test_pbt_auto_download.py
class TestMissingCheckpointTriggersDownload (line 66) | class TestMissingCheckpointTriggersDownload:
method test_missing_pth_triggers_download_and_returns_valid_path (line 79) | def test_missing_pth_triggers_download_and_returns_valid_path(
class TestExistingCheckpointSkipsDownload (line 141) | class TestExistingCheckpointSkipsDownload:
method test_existing_pth_skips_download (line 153) | def test_existing_pth_skips_download(self, pth_name: str) -> None:
class TestAutoDownloadIsTorchOnly (line 186) | class TestAutoDownloadIsTorchOnly:
method test_non_pth_extension_raises_without_download (line 198) | def test_non_pth_extension_raises_without_download(self, ext: str) -> ...
function _make_hf_hub_http_error (line 225) | def _make_hf_hub_http_error(message: str) -> Exception:
class TestNetworkErrorsProduceActionableMessages (line 253) | class TestNetworkErrorsProduceActionableMessages:
method test_network_errors_produce_actionable_messages (line 266) | def test_network_errors_produce_actionable_messages(
FILE: tests/test_pbt_backend_resolution.py
function _expected_backend (line 30) | def _expected_backend(cli: str | None, env: str | None, auto: str) -> str:
function test_backend_resolution_priority_chain (line 46) | def test_backend_resolution_priority_chain(cli: str | None, env: str | N...
FILE: tests/test_pbt_dep_preservation.py
function _parse_base_dependencies (line 53) | def _parse_base_dependencies() -> list[str]:
function _normalize_dep (line 60) | def _normalize_dep(dep: str) -> str:
function _dep_present (line 69) | def _dep_present(dep: str, dep_list: list[str]) -> bool:
function test_non_torch_dependency_preserved (line 80) | def test_non_torch_dependency_preserved(dep: str) -> None:
FILE: tests/test_pyproject_structure.py
function pyproject (line 26) | def pyproject() -> dict:
class TestCudaExtra (line 37) | class TestCudaExtra:
method test_cuda_extra_contains_torch (line 40) | def test_cuda_extra_contains_torch(self, pyproject: dict) -> None:
method test_cuda_extra_contains_torchvision (line 44) | def test_cuda_extra_contains_torchvision(self, pyproject: dict) -> None:
class TestMlxExtra (line 54) | class TestMlxExtra:
method test_mlx_extra_contains_corridorkey_mlx (line 57) | def test_mlx_extra_contains_corridorkey_mlx(self, pyproject: dict) -> ...
class TestPytorchIndex (line 68) | class TestPytorchIndex:
method test_pytorch_index_has_cuda_extra (line 71) | def test_pytorch_index_has_cuda_extra(self, pyproject: dict) -> None:
class TestUvSources (line 83) | class TestUvSources:
method test_torch_source_has_cuda_extra (line 86) | def test_torch_source_has_cuda_extra(self, pyproject: dict) -> None:
method test_torchvision_source_has_cuda_extra (line 92) | def test_torchvision_source_has_cuda_extra(self, pyproject: dict) -> N...
class TestConflicts (line 104) | class TestConflicts:
method test_cuda_mlx_conflict_declared (line 107) | def test_cuda_mlx_conflict_declared(self, pyproject: dict) -> None:
class TestTimmSourcePreserved (line 121) | class TestTimmSourcePreserved:
method test_timm_source_is_git (line 124) | def test_timm_source_is_git(self, pyproject: dict) -> None:
method test_timm_git_url (line 128) | def test_timm_git_url(self, pyproject: dict) -> None:
method test_timm_git_branch (line 132) | def test_timm_git_branch(self, pyproject: dict) -> None:
class TestTritonWindowsPreserved (line 142) | class TestTritonWindowsPreserved:
method test_triton_windows_in_base_deps (line 145) | def test_triton_windows_in_base_deps(self, pyproject: dict) -> None:
method test_triton_windows_has_win32_marker (line 150) | def test_triton_windows_has_win32_marker(self, pyproject: dict) -> None:
class TestDevDependencyGroup (line 161) | class TestDevDependencyGroup:
method test_dev_group_contains_pytest (line 164) | def test_dev_group_contains_pytest(self, pyproject: dict) -> None:
method test_dev_group_contains_pytest_cov (line 168) | def test_dev_group_contains_pytest_cov(self, pyproject: dict) -> None:
method test_dev_group_contains_ruff (line 172) | def test_dev_group_contains_ruff(self, pyproject: dict) -> None:
class TestDocsDependencyGroup (line 182) | class TestDocsDependencyGroup:
method test_docs_group_contains_zensical (line 185) | def test_docs_group_contains_zensical(self, pyproject: dict) -> None:
Condensed preview — 92 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (627K chars).
[
{
"path": ".dockerignore",
"chars": 522,
"preview": ".git\n.github\n.pytest_cache\n.ruff_cache\n.uv-cache\n.venv\n__pycache__/\n*.pyc\n*.pyo\n*.pyd\n\n# Docker files (no need to copy t"
},
{
"path": ".git-blame-ignore-revs",
"chars": 120,
"preview": "# Automated code formatting — no behavioral changes\n# ruff format + lint fixes\nb0ad00efbc791ed097cd3fd241c10319beb8a631\n"
},
{
"path": ".github/ISSUE_TEMPLATE/bug_report.yml",
"chars": 3306,
"preview": "name: 🐛 Bug Report\ndescription: Report a reproducible issue or unexpected behavior in the project.\ntitle: \"[Bug]: \"\nlabe"
},
{
"path": ".github/pull_request_template.md",
"chars": 168,
"preview": "## What does this change?\n\n## How was it tested?\n\n## Checklist\n\n- [ ] `uv run pytest` passes\n- [ ] `uv run ruff check` p"
},
{
"path": ".github/workflows/ci.yml",
"chars": 1150,
"preview": "name: CI\n\non:\n push:\n branches: [main]\n pull_request:\n branches: [main]\n\nenv:\n UV_NO_SYNC: 1\n UV_LOCKED: 1\n O"
},
{
"path": ".github/workflows/docs.yml",
"chars": 1015,
"preview": "name: Documentation\non:\n push:\n branches:\n - main\n paths:\n - \"docs/**\"\n "
},
{
"path": ".gitignore",
"chars": 798,
"preview": "# Python\n__pycache__/\n*.pyc\n*.pyo\n*.pyd\n.coverage\n.Python\n.pytest_cache/\n.venv/\n.hypothesis\nenv/\nvenv/\nENV/\nenv.bak/\nven"
},
{
"path": ".python-version",
"chars": 4,
"preview": "3.13"
},
{
"path": "BiRefNetModule/checkpoints/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "BiRefNetModule/wrapper.py",
"chars": 7399,
"preview": "import logging\nimport os\nfrom pathlib import Path\nfrom typing import Tuple\n\nimport cv2\nimport numpy as np\nimport torch\nf"
},
{
"path": "CONTRIBUTING.md",
"chars": 4551,
"preview": "# Contributing to CorridorKey\n\nThanks for your interest in improving CorridorKey! Whether you're a VFX artist, a pipelin"
},
{
"path": "ClipsForInference/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "CorridorKeyModule/IgnoredCheckpoints/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "CorridorKeyModule/README.md",
"chars": 3646,
"preview": "# CorridorKeyModule\n\nA self-contained, high-performance AI Chroma Keying engine. This module provides a simple API to ac"
},
{
"path": "CorridorKeyModule/__init__.py",
"chars": 105,
"preview": "from __future__ import annotations\n\nfrom .inference_engine import CorridorKeyEngine as CorridorKeyEngine\n"
},
{
"path": "CorridorKeyModule/backend.py",
"chars": 11264,
"preview": "\"\"\"Backend factory — selects Torch or MLX engine and normalizes output contracts.\"\"\"\n\nfrom __future__ import annotations"
},
{
"path": "CorridorKeyModule/checkpoints/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "CorridorKeyModule/core/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "CorridorKeyModule/core/color_utils.py",
"chars": 9731,
"preview": "from __future__ import annotations\n\nimport functools\nfrom collections.abc import Callable\n\nimport cv2\nimport numpy as np"
},
{
"path": "CorridorKeyModule/core/model_transformer.py",
"chars": 11271,
"preview": "from __future__ import annotations\n\nimport logging\n\nimport timm\nimport torch\nimport torch.nn as nn\nimport torch.nn.funct"
},
{
"path": "CorridorKeyModule/inference_engine.py",
"chars": 10806,
"preview": "from __future__ import annotations\n\nimport logging\nimport math\nimport os\nimport sys\n\nimport cv2\nimport numpy as np\nimpor"
},
{
"path": "CorridorKey_DRAG_CLIPS_HERE_local.bat",
"chars": 898,
"preview": "@echo off\nREM Corridor Key Launcher - Local\n\nREM Set script path (assumes corridorkey_cli.py is in the same directory as"
},
{
"path": "CorridorKey_DRAG_CLIPS_HERE_local.sh",
"chars": 1528,
"preview": "#!/usr/bin/env bash\n# Corridor Key Launcher - Local Linux/macOS\n\ncd \"$(dirname \"$0\")\"\n\n# Get the directory where this sc"
},
{
"path": "Dockerfile",
"chars": 1075,
"preview": "FROM ghcr.io/astral-sh/uv:0.7-python3.11-bookworm-slim\n\n# Create non-root user upfront.\nRUN useradd --create-home --uid "
},
{
"path": "IgnoredClips/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "Install_CorridorKey_Linux_Mac.sh",
"chars": 2634,
"preview": "#!/usr/bin/env bash\n\ncd \"$(dirname \"$0\")\"\n\necho \"===================================================\"\necho \" Corridor"
},
{
"path": "Install_CorridorKey_Windows.bat",
"chars": 2166,
"preview": "@echo off\nTITLE CorridorKey Setup Wizard\necho ===================================================\necho CorridorKey -"
},
{
"path": "Install_GVM_Linux_Mac.sh",
"chars": 1072,
"preview": "#!/usr/bin/env bash\n\ncd \"$(dirname \"$0\")\"\n\n# Set the Terminal window title\necho -n -e \"\\033]0;GVM Setup Wizard\\007\"\necho"
},
{
"path": "Install_GVM_Windows.bat",
"chars": 906,
"preview": "@echo off\nTITLE GVM Setup Wizard\necho ===================================================\necho GVM (AlphaHint Genera"
},
{
"path": "Install_VideoMaMa_Linux_Mac.sh",
"chars": 1112,
"preview": "#!/usr/bin/env bash\n\ncd \"$(dirname \"$0\")\"\n\n# Set the Terminal window title\necho -n -e \"\\033]0;VideoMaMa Setup Wizard\\007"
},
{
"path": "Install_VideoMaMa_Windows.bat",
"chars": 965,
"preview": "@echo off\nTITLE VideoMaMa Setup Wizard\necho ===================================================\necho VideoMaMa (AlphaH"
},
{
"path": "LICENSE",
"chars": 24771,
"preview": "CORRIDOR KEY LICENCE\n=======================================================================\n\nVersion 1.0\n\nCopyright (c)"
},
{
"path": "Output/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "README.md",
"chars": 23025,
"preview": "# CorridorKey\n\n\nhttps://github.com/user-attachments/assets/1fb27ea8-bc91-4ebc-818f-5a3b5585af08\n\n\nWhen you film somethin"
},
{
"path": "RunGVMOnly.sh",
"chars": 410,
"preview": "#!/usr/bin/env bash\n\n# Ensure script stops on error\nset -e\n\n# Path to script directory\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BA"
},
{
"path": "RunInferenceOnly.sh",
"chars": 452,
"preview": "#!/usr/bin/env bash\n\n# Ensure script stops on error\nset -e\n\n# Path to script directory\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BA"
},
{
"path": "VideoMaMaInferenceModule/LICENSE.md",
"chars": 1053,
"preview": "# VideoMaMa Licensing and Acknowledgements\n\nThis module (`VideoMaMaInferenceModule`) contains repackaged code and integr"
},
{
"path": "VideoMaMaInferenceModule/README.md",
"chars": 1409,
"preview": "# VideoMaMa Inference Module\n\nThis module provides a standalone interface for running VideoMaMa inference.\n\n## Usage\n\n``"
},
{
"path": "VideoMaMaInferenceModule/__init__.py",
"chars": 287,
"preview": "from .inference import load_videomama_model, run_inference, extract_frames_from_video, save_video\nfrom .pipeline import "
},
{
"path": "VideoMaMaInferenceModule/checkpoints/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "VideoMaMaInferenceModule/inference.py",
"chars": 7457,
"preview": "\"\"\"\nVideoMaMa Inference Module\nProvides functions to load the model and run inference on video inputs.\n\"\"\"\n\nimport os\nim"
},
{
"path": "VideoMaMaInferenceModule/pipeline.py",
"chars": 48646,
"preview": "# pipeline_svd_masked.py\n\nimport inspect\nfrom dataclasses import dataclass\nfrom typing import Callable, Dict, List, Opti"
},
{
"path": "backend/__init__.py",
"chars": 1526,
"preview": "\"\"\"Backend service layer for ez-CorridorKey.\"\"\"\n\nfrom .clip_state import (\n ClipAsset,\n ClipEntry,\n ClipState,\n"
},
{
"path": "backend/clip_state.py",
"chars": 18826,
"preview": "\"\"\"Clip entry data model and state machine.\n\nState Machine:\n EXTRACTING — Video input being extracted to image sequen"
},
{
"path": "backend/errors.py",
"chars": 3917,
"preview": "\"\"\"Typed exceptions for the CorridorKey backend.\"\"\"\n\nimport sys\n\n\nclass CorridorKeyError(Exception):\n \"\"\"Base excepti"
},
{
"path": "backend/ffmpeg_tools.py",
"chars": 12357,
"preview": "\"\"\"FFmpeg subprocess wrapper for video extraction and stitching.\n\nPure Python, no Qt deps. Provides:\n- find_ffmpeg() / f"
},
{
"path": "backend/frame_io.py",
"chars": 5960,
"preview": "\"\"\"Unified frame I/O — read images and video frames as float32 RGB.\n\nAll reading functions return float32 arrays in [0, "
},
{
"path": "backend/job_queue.py",
"chars": 12309,
"preview": "\"\"\"GPU job queue with mutual exclusion.\n\nEnsures only ONE GPU job runs at a time across all job types\n(inference, GVM al"
},
{
"path": "backend/natural_sort.py",
"chars": 879,
"preview": "\"\"\"Natural sort key for frame filenames.\n\nHandles non-zero-padded frame numbers correctly:\n frame_1, frame_2, frame_10 "
},
{
"path": "backend/project.py",
"chars": 12728,
"preview": "\"\"\"Project folder management — creation, scanning, and metadata.\n\nA project is a timestamped container holding one or mo"
},
{
"path": "backend/service.py",
"chars": 40867,
"preview": "\"\"\"CorridorKeyService — clean backend API for the UI and CLI.\n\nThis module wraps all processing logic from clip_manager."
},
{
"path": "backend/validators.py",
"chars": 4335,
"preview": "\"\"\"Validation utilities for frame processing.\n\nAll validators either return cleaned data or raise typed exceptions from "
},
{
"path": "clip_manager.py",
"chars": 39063,
"preview": "from __future__ import annotations\n\nimport argparse\nimport glob\nimport logging\nimport os\nimport shutil\nimport sys\nfrom d"
},
{
"path": "corridorkey_cli.py",
"chars": 21310,
"preview": "\"\"\"CorridorKey command-line interface and interactive wizard.\n\nThis module handles CLI subcommands, environment setup, a"
},
{
"path": "device_utils.py",
"chars": 2473,
"preview": "\"\"\"Centralized cross-platform device selection for CorridorKey.\"\"\"\n\nimport logging\nimport os\n\nimport torch\n\nlogger = log"
},
{
"path": "docker-compose.yml",
"chars": 1300,
"preview": "services:\n corridorkey:\n profiles: [\"gpu\"]\n build:\n context: .\n dockerfile: Dockerfile\n image: corri"
},
{
"path": "docs/LLM_HANDOVER.md",
"chars": 5821,
"preview": "# CorridorKey: LLM Handover Guide\n\nWelcome, fellow AI Assistant! You are picking up a highly specialized computer vision"
},
{
"path": "docs/index.md",
"chars": 0,
"preview": ""
},
{
"path": "gvm_core/LICENSE.md",
"chars": 852,
"preview": "# GVM Licensing and Acknowledgements\n\nThis module (`gvm_core`) contains repackaged code and integrations from the **Gene"
},
{
"path": "gvm_core/README.md",
"chars": 2413,
"preview": "# GVM Core Module\n\nThis folder contains the core logic and pre-trained models for Generative Video Matting (GVM).\nIt is "
},
{
"path": "gvm_core/__init__.py",
"chars": 34,
"preview": "from .wrapper import GVMProcessor\n"
},
{
"path": "gvm_core/gvm/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "gvm_core/gvm/models/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "gvm_core/gvm/models/unet_spatio_temporal_condition.py",
"chars": 27343,
"preview": "from dataclasses import dataclass\nfrom typing import Dict, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\n\nf"
},
{
"path": "gvm_core/gvm/pipelines/pipeline_gvm.py",
"chars": 8913,
"preview": "import torch\nimport tqdm\nimport numpy as np\nfrom diffusers import DiffusionPipeline\nfrom diffusers.utils import (\n Ba"
},
{
"path": "gvm_core/gvm/utils/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "gvm_core/gvm/utils/inference_utils.py",
"chars": 5628,
"preview": "import av\nimport os\nimport pims\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset\nfrom torchvision.tr"
},
{
"path": "gvm_core/weights/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "gvm_core/wrapper.py",
"chars": 11130,
"preview": "import os\nimport os.path as osp\nimport cv2\nimport random\nimport logging\nimport time\nfrom pathlib import Path\n\nfrom easyd"
},
{
"path": "pyproject.toml",
"chars": 3627,
"preview": "[project]\nname = \"corridorkey\"\nversion = \"1.0.0\"\ndescription = \"Neural network green screen keying for professional VFX "
},
{
"path": "renovate.json",
"chars": 257,
"preview": "{\n \"$schema\": \"https://docs.renovatebot.com/renovate-schema.json\",\n \"extends\": [\"config:recommended\"],\n \"pinDigests\":"
},
{
"path": "test_vram.py",
"chars": 987,
"preview": "import timeit\n\nimport numpy as np\nimport torch\n\nfrom CorridorKeyModule.inference_engine import CorridorKeyEngine\n\n\ndef p"
},
{
"path": "tests/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "tests/conftest.py",
"chars": 7578,
"preview": "\"\"\"Shared pytest configuration and fixtures for CorridorKey tests.\"\"\"\n\nimport platform\nimport sys\nfrom types import Modu"
},
{
"path": "tests/test_backend.py",
"chars": 11259,
"preview": "\"\"\"Unit tests for CorridorKeyModule.backend — no GPU/MLX required.\"\"\"\n\nimport errno\nimport logging\nimport os\nfrom unitte"
},
{
"path": "tests/test_cli.py",
"chars": 8448,
"preview": "\"\"\"Tests for the typer-based CLI in corridorkey_cli.py.\"\"\"\n\nfrom __future__ import annotations\n\nimport re\nfrom unittest."
},
{
"path": "tests/test_clip_manager.py",
"chars": 35391,
"preview": "\"\"\"Tests for clip_manager.py utility functions and ClipEntry discovery.\n\nThese tests verify the non-interactive parts of"
},
{
"path": "tests/test_color_utils.py",
"chars": 22678,
"preview": "\"\"\"Unit tests for CorridorKeyModule.core.color_utils.\n\nThese tests verify the color math that underpins CorridorKey's co"
},
{
"path": "tests/test_device_utils.py",
"chars": 6091,
"preview": "\"\"\"Unit tests for device_utils — cross-platform device selection.\n\nTests cover all code paths in detect_best_device(), r"
},
{
"path": "tests/test_e2e_workflow.py",
"chars": 4855,
"preview": "\"\"\"End-to-end workflow integration tests for CorridorKey.\n\nThese tests exercise the full pipeline from ClipEntry asset d"
},
{
"path": "tests/test_exr_gamma_bug_condition.py",
"chars": 11024,
"preview": "\"\"\"Bug condition exploration tests for EXR gamma correction.\n\nThese tests encode the EXPECTED (correct) behavior for EXR"
},
{
"path": "tests/test_exr_gamma_preservation.py",
"chars": 12289,
"preview": "\"\"\"Preservation property tests for EXR gamma correction bugfix.\n\nThese tests verify that non-buggy code paths remain unc"
},
{
"path": "tests/test_frame_io.py",
"chars": 2272,
"preview": "\"\"\"Tests for backend.frame_io — frame reading utilities.\n\nFocuses on input validation and edge cases that don't require "
},
{
"path": "tests/test_gamma_consistency.py",
"chars": 5833,
"preview": "\"\"\"Tests documenting the gamma 2.2 vs piecewise sRGB inconsistency.\n\nSTATUS: This test documents a KNOWN INCONSISTENCY, "
},
{
"path": "tests/test_imports.py",
"chars": 917,
"preview": "\"\"\"Smoke tests: verify all CorridorKey packages import without error.\n\nThese catch missing __init__.py files, broken rel"
},
{
"path": "tests/test_inference_engine.py",
"chars": 12922,
"preview": "\"\"\"Tests for CorridorKeyModule.inference_engine.CorridorKeyEngine.process_frame.\n\nThese tests mock the GreenFormer model"
},
{
"path": "tests/test_mlx_smoke.py",
"chars": 1684,
"preview": "\"\"\"MLX integration smoke test — requires Apple Silicon + corridorkey_mlx.\"\"\"\n\nimport numpy as np\nimport pytest\n\npytestma"
},
{
"path": "tests/test_pbt_auto_download.py",
"chars": 10812,
"preview": "\"\"\"Feature: auto-model-download — Property-based tests for auto checkpoint download.\n\nProperties tested:\n 1: Missing ch"
},
{
"path": "tests/test_pbt_backend_resolution.py",
"chars": 2562,
"preview": "\"\"\"Property-based test for backend resolution priority chain.\n\n# Feature: uv-lock-drift-fix, Property 1: Backend resolut"
},
{
"path": "tests/test_pbt_dep_preservation.py",
"chars": 2747,
"preview": "\"\"\"Property-based test for non-torch dependency preservation.\n\n# Feature: uv-lock-drift-fix, Property 2: Non-torch depen"
},
{
"path": "tests/test_pyproject_structure.py",
"chars": 7383,
"preview": "\"\"\"Structural validation tests for pyproject.toml extras configuration.\n\nValidates that the pyproject.toml correctly def"
},
{
"path": "zensical.toml",
"chars": 1335,
"preview": "[project]\nsite_name = \"CorridorKey\"\nsite_description = \"Perfect Green Screen Keys\"\nrepo_name = \"CorridorKey\"\nrepo_url = "
}
]
About this extraction
This page contains the full source code of the nikopueringer/CorridorKey GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 92 files (584.6 KB), approximately 141.7k tokens, and a symbol index with 655 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.