reshot 0.3.7__py3-none-any.whl

reshot/__init__.py

@@ -0,0 +1,58 @@
+"""ReShot — turn any video into a depth-map video for video-generation control.
+
+from pathlib import Path
+from reshot import RunConfig, run
+run(RunConfig(input=Path("in.mp4"), output=Path("out.mp4"), target="h3"))
+
+from reshot import extract, to_gray, write_gray_video
+depths, fps = extract("in.mp4")            # float32 [T, H, W], larger = closer
+write_gray_video(to_gray(depths), "out.mp4", fps)
+"""
+
+from __future__ import annotations
+
+import os
+
+# Apple Silicon: cap the MPS allocator at 60 % of the recommended working set. Unified
+# memory means GPU allocations count against the same RAM as everything else; without a
+# cap a long clip grows until the OS kills processes (seen 2026-09-11). Must be set
+# before torch is imported, which is why it lives here.
+os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.6")
+
+from ._version import __version__
+from .config import RunConfig
+from .errors import BackendError, InputError, RamBudgetError, ReshotError, ToolMissingError
+from .io import probe_video, read_video, write_gray_video
+from .pipeline import Plan, RunResult, extract, model_input_resolution, plan, resolve_input_size, run, run_many
+from .planning import memory_verdict, processing_max_res
+from .postprocess import to_gray, upsample_frames
+from .targets import TARGETS, Target, center_crop, fit_dimensions
+
+__all__ = [
+    "TARGETS",
+    "BackendError",
+    "InputError",
+    "Plan",
+    "RamBudgetError",
+    "ReshotError",
+    "RunConfig",
+    "RunResult",
+    "Target",
+    "ToolMissingError",
+    "__version__",
+    "center_crop",
+    "extract",
+    "fit_dimensions",
+    "memory_verdict",
+    "model_input_resolution",
+    "plan",
+    "probe_video",
+    "processing_max_res",
+    "read_video",
+    "resolve_input_size",
+    "run",
+    "run_many",
+    "to_gray",
+    "upsample_frames",
+    "write_gray_video",
+]

reshot/_version.py

@@ -0,0 +1 @@
+__version__ = "0.3.7"

reshot/backends/__init__.py

@@ -0,0 +1,23 @@
+"""Backend registry. `get_backend("vda", model="small", device="auto")`."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .base import BackendInfo, DepthBackend
+from .fake import FakeBackend
+from .vda import VDABackend, pick_device
+
+
+def get_backend(name: str, **kwargs: Any) -> DepthBackend:
+    """Instantiate a backend by name. Unknown names raise `BackendError`."""
+    from ..errors import BackendError
+
+    if name == "vda":
+        return VDABackend(**kwargs)
+    if name == "fake":
+        return FakeBackend()
+    raise BackendError(f"unknown backend {name!r}")
+
+
+__all__ = ["BackendInfo", "DepthBackend", "FakeBackend", "VDABackend", "get_backend", "pick_device"]

reshot/backends/base.py

@@ -0,0 +1,34 @@
+"""Backend protocol: a depth model that turns a frame stack into per-frame depth.
+
+Every backend returns *relative inverse depth* (larger = closer to camera) as float32
+`[T, H, W]` at the input frame resolution, already temporally aligned across the whole
+clip. Normalisation to 8-bit grey happens later in :mod:`reshot.postprocess`,
+so backends never decide what "white" means.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class BackendInfo:
+    name: str
+    variant: str
+    license: str
+    commercial_ok: bool
+
+
+class DepthBackend(Protocol):
+    info: BackendInfo
+
+    def infer(self, frames: np.ndarray, fps: float, *, input_size: int = 518) -> np.ndarray:
+        """`frames`: uint8 RGB `[T, H, W, 3]` → float32 inverse depth `[T, H, W]`."""
+        ...
+
+    def peak_memory_bytes(self) -> dict[str, int]:
+        """Accelerator memory peak of the last `infer()`; `{}` when there is none to report."""
+        ...

reshot/backends/fake.py

@@ -0,0 +1,28 @@
+"""A backend that returns a plausible depth field without loading a model.
+
+For exercising everything *around* inference — I/O, encoding, presets, metrics,
+remote job runners — in CI or on a laptop, in milliseconds and without downloads.
+Select with `--backend fake` (or the env var `RESHOT_FAKE_BACKEND=1`)."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .base import BackendInfo
+
+
+class FakeBackend:
+    info = BackendInfo(name="fake", variant="none", license="n/a", commercial_ok=True)
+    fp32 = True
+    device = "cpu"
+
+    def infer(self, frames: np.ndarray, fps: float, *, input_size: int = 518) -> np.ndarray:
+        t, h, w = frames.shape[:3]
+        yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
+        # a soft radial "subject" in the middle that brightens over time
+        base = 1.0 / (1.0 + ((xx - w / 2) ** 2 + (yy - h / 2) ** 2) / (0.15 * w * h))
+        ramp = np.linspace(0.8, 1.0, num=max(t, 1), dtype=np.float32)
+        return (base[None] * ramp[:, None, None]).astype(np.float32)
+
+    def peak_memory_bytes(self) -> dict[str, int]:
+        return {}

reshot/backends/vda.py

@@ -0,0 +1,140 @@
+"""Video Depth Anything backend (default).
+
+Why this model: it is the only open video-depth model that ships a temporal module
+*and* an Apache-2.0 checkpoint (Small). Per-frame models (Depth Anything V2/3) flicker
+on video unless you bolt on your own smoothing; diffusion models (DepthCrafter) are
+~100× slower. VDA-Small at 28M params runs on an 8 GB consumer GPU or an Apple M-series.
+
+Precision policy (measured 2026-09-11 on M2 Max, 32 frames @ 736×1280):
+    mps  fp32 →  7.0 s        mps fp16 → did not finish in 3 min (autocast pathological)
+    cpu  fp32 → 59 s
+    Full 294-frame clip end-to-end on mps fp32: 147 s (~500 ms/frame incl. alignment).
+So fp16 is enabled on CUDA only. Do not "optimise" this by turning fp16 on for MPS.
+
+Licence policy: `small` is Apache-2.0. `base`/`large` are CC-BY-NC-4.0 — loading them
+prints a warning and sets ``info.commercial_ok = False``; the CLI surfaces this.
+"""
+
+from __future__ import annotations
+
+import logging
+import warnings
+
+import numpy as np
+import torch
+
+from .base import BackendInfo
+
+log = logging.getLogger(__name__)
+
+_VARIANTS = {
+    "small": {"encoder": "vits", "features": 64, "out_channels": [48, 96, 192, 384]},
+    "base": {"encoder": "vitb", "features": 128, "out_channels": [96, 192, 384, 768]},
+    "large": {"encoder": "vitl", "features": 256, "out_channels": [256, 512, 1024, 1024]},
+}
+_LICENSES = {"small": "Apache-2.0", "base": "CC-BY-NC-4.0", "large": "CC-BY-NC-4.0"}
+_HF_REPOS = {
+    "small": "depth-anything/Video-Depth-Anything-Small",
+    "base": "depth-anything/Video-Depth-Anything-Base",
+    "large": "depth-anything/Video-Depth-Anything-Large",
+}
+
+
+def pick_device(requested: str = "auto") -> str:
+    if requested != "auto":
+        return requested
+    if torch.cuda.is_available():
+        return "cuda"
+    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def _weights_path(variant: str, checkpoint: str | None) -> str:
+    if checkpoint:
+        return checkpoint
+    # huggingface_hub honours HF_ENDPOINT, so users in China can point it at a mirror
+    # (e.g. https://hf-mirror.com) without us special-casing anything.
+    from huggingface_hub import hf_hub_download
+
+    filename = f"video_depth_anything_{_VARIANTS[variant]['encoder']}.pth"
+    return hf_hub_download(_HF_REPOS[variant], filename)
+
+
+class VDABackend:
+    def __init__(
+        self,
+        model: str = "small",
+        device: str = "auto",
+        checkpoint: str | None = None,
+        cuda_memory_fraction: float | None = None,
+    ) -> None:
+        variant = model
+        if variant not in _VARIANTS:
+            from ..errors import BackendError
+
+            raise BackendError(f"unknown model variant {variant!r}; choose from {sorted(_VARIANTS)}")
+        from ..third_party.video_depth_anything.video_depth import VideoDepthAnything
+
+        self.variant = variant
+        self.device = pick_device(device)
+        # fp16 only where it is known to be fast; see module docstring for the measurement.
+        self.fp32 = self.device != "cuda"
+        self.info = BackendInfo(
+            name="video-depth-anything",
+            variant=variant,
+            license=_LICENSES[variant],
+            commercial_ok=_LICENSES[variant].startswith("Apache"),
+        )
+        if not self.info.commercial_ok:
+            warnings.warn(
+                f"Video Depth Anything '{variant}' weights are {_LICENSES[variant]}: "
+                "non-commercial use only. Use 'small' for commercial projects.",
+                stacklevel=2,
+            )
+
+        try:
+            path = _weights_path(variant, checkpoint)
+        except Exception as exc:  # hub / network / offline-cache errors all land here
+            from ..errors import BackendError
+
+            raise BackendError(
+                f"could not obtain weights for '{variant}': {exc}. "
+                "Set HF_ENDPOINT to a mirror, or pass --checkpoint with a local .pth."
+            ) from exc
+        if cuda_memory_fraction is not None and self.device == "cuda":
+            # Emulate a smaller card: allocations beyond the fraction raise OOM instead of
+            # silently using the whole GPU. Used to verify the "runs on N GB" claim.
+            torch.cuda.set_per_process_memory_fraction(cuda_memory_fraction)
+            log.info("cuda memory capped at %.0f%% of the card", cuda_memory_fraction * 100)
+        log.info("loading %s from %s on %s (fp32=%s)", variant, path, self.device, self.fp32)
+        model = VideoDepthAnything(**_VARIANTS[variant])
+        model.load_state_dict(torch.load(path, map_location="cpu", weights_only=True), strict=True)
+        self.model = model.to(self.device).eval()
+
+    def infer(self, frames: np.ndarray, fps: float, *, input_size: int = 518) -> np.ndarray:
+        if frames.ndim != 4 or frames.shape[-1] != 3:
+            raise ValueError(f"expected [T, H, W, 3] uint8 RGB, got {frames.shape}")
+        # Upstream's infer_video_depth already does 32-frame windows with 10-frame overlap,
+        # keyframe-based scale/shift alignment across windows and interpolation over the
+        # seam. That is the temporal-consistency machinery — do not re-chunk outside it.
+        if self.device == "cuda":
+            torch.cuda.reset_peak_memory_stats()
+        with torch.inference_mode():  # upstream uses no_grad; inference_mode also skips version counters
+            depths, _ = self.model.infer_video_depth(
+                frames, fps, input_size=input_size, device=self.device, fp32=self.fp32
+            )
+        return np.asarray(depths, dtype=np.float32)
+
+    def peak_memory_bytes(self) -> dict[str, int]:
+        """Peak VRAM of the last `infer()`: what the model actually touched (`allocated`)
+        and what the caching allocator held (`reserved`, ≈ what nvidia-smi shows)."""
+        if self.device != "cuda":
+            return {}
+        return {
+            "gpu_peak_allocated_bytes": int(torch.cuda.max_memory_allocated()),
+            "gpu_peak_reserved_bytes": int(torch.cuda.max_memory_reserved()),
+        }
+
+
+__all__ = ["VDABackend", "pick_device"]

reshot/cli.py

@@ -0,0 +1,169 @@
+"""`reshot` command line: argv → RunConfig(s) → pipeline.run() / run_many().
+
+    reshot in.mp4 -o out.mp4                 # Apache-2.0 Small model, auto device
+    reshot in.mp4 -o out.mp4 --target h3     # 24 fps, ×32 dims, ≤15 s for MiniMax H3
+    reshot in.mp4 -o out.mp4 --npz d.npz --metrics run.json
+    reshot clips/*.mp4 -o depth/ --target seedance   # batch: one model load, out/<name>_depth.mp4
+
+Batch mode (several inputs, or `-o` naming a directory): `--npz` / `--metrics`, if given,
+are directories too and get `<name>.npz` / `<name>.json` per clip. One bad clip is
+reported and skipped; the exit code is that of the first failure.
+
+Exit codes: 0 ok · 1 unexpected · 2 bad input/arguments · 3 over RAM budget ·
+4 ffmpeg missing · 5 backend/weights problem.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+from ._version import __version__
+from .config import BACKENDS, MODEL_VARIANTS, QUALITIES, RunConfig
+from .errors import ReshotError
+from .pipeline import run, run_many
+from .reporter import StderrReporter
+from .targets import TARGETS
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="reshot",
+        description="Video → depth-map video for video-generation ControlNets.",
+        epilog="Docs: https://github.com/maosika-ai/reshot",
+    )
+    p.add_argument("input", type=Path, nargs="+", help="input video(s) (anything ffmpeg/OpenCV can read)")
+    p.add_argument(
+        "-o", "--output", type=Path, required=True, help="output .mp4, or a directory when there are several inputs"
+    )
+    g = p.add_argument_group("model")
+    g.add_argument(
+        "--model",
+        default="small",
+        choices=MODEL_VARIANTS,
+        help="Video Depth Anything variant; only 'small' is Apache-2.0 (default)",
+    )
+    g.add_argument("--backend", default="vda", choices=BACKENDS, help=argparse.SUPPRESS)
+    g.add_argument("--device", default="auto", help="auto | cuda | mps | cpu")
+    g.add_argument(
+        "--quality",
+        default="fast",
+        choices=QUALITIES,
+        help="what the model sees: fast = 644x364 for 16:9 / 364x644 for 9:16, ~3 GB VRAM (default); "
+        "full = 924x518 / 518x924, ~11 GB VRAM, 2.5x slower, sharper fine detail",
+    )
+    g.add_argument(
+        "--input-size",
+        type=int,
+        default=None,
+        help="expert: model short side in px, multiple of 14; overrides --quality",
+    )
+    g.add_argument("--checkpoint", type=Path, help="local .pth instead of the Hugging Face download")
+    g = p.add_argument_group("output")
+    g.add_argument(
+        "--target",
+        default="none",
+        choices=sorted(TARGETS),
+        help="fps / frame-size preset for a generator: " + ", ".join(f"{k}={v.note}" for k, v in TARGETS.items()),
+    )
+    g.add_argument("--fps", type=float, help="override output fps (default: preset or source)")
+    g.add_argument(
+        "--max-res", type=int, default=1280, help="cap the OUTPUT's longer side (inference runs at model resolution)"
+    )
+    g.add_argument("--max-frames", type=int, help="stop after N frames")
+    g.add_argument("--invert", action="store_true", help="far = white instead of near = white")
+    g.add_argument(
+        "--clip", type=float, default=0.0, metavar="PCT", help="percentile clip on both tails before scaling"
+    )
+    g.add_argument("--gamma", type=float, default=1.0, help=">1 darkens mid-tones")
+    g.add_argument("--crf", type=int, default=12, help="x264 CRF; lower = larger, cleaner (default 12)")
+    g.add_argument("--npz", type=Path, help="also save raw float depth (processing resolution)")
+    g.add_argument("--metrics", type=Path, help="write timing / memory / size metrics as JSON")
+    g = p.add_argument_group("safety")
+    g.add_argument("--force", action="store_true", help="run even if the RAM estimate exceeds the budget")
+    # Verification aid, not a user feature: makes a big card behave like a small one so a
+    # "runs on N GB" claim can be tested (torch.cuda.set_per_process_memory_fraction).
+    g.add_argument("--cuda-memory-fraction", type=float, default=None, help=argparse.SUPPRESS)
+    p.add_argument("-v", "--verbose", action="store_true")
+    p.add_argument("--version", action="version", version=f"reshot {__version__}")
+    return p
+
+
+def _is_batch(args: argparse.Namespace) -> bool:
+    """Several inputs, or an output that is (or is spelled like) a directory."""
+    o = args.output
+    return len(args.input) > 1 or o.is_dir() or str(o).endswith(("/", "\\"))
+
+
+def configs_from_args(args: argparse.Namespace) -> list[RunConfig]:
+    """One RunConfig per input. In batch mode the outputs are derived from the input names."""
+    batch = _is_batch(args)
+
+    def per_clip(option: Path | None, stem: str, ext: str) -> Path | None:
+        if option is None:
+            return None
+        return option / f"{stem}{ext}" if batch else option
+
+    cfgs = []
+    for src in args.input:
+        stem = src.stem
+        cfgs.append(
+            RunConfig(
+                input=src,
+                output=args.output / f"{stem}_depth.mp4" if batch else args.output,
+                model=args.model,
+                backend=args.backend,
+                device=args.device,
+                target=args.target,
+                fps=args.fps,
+                max_res=args.max_res,
+                max_frames=args.max_frames,
+                quality=args.quality,
+                input_size=args.input_size,
+                invert=args.invert,
+                clip_percent=args.clip,
+                gamma=args.gamma,
+                crf=args.crf,
+                npz=per_clip(args.npz, stem, ".npz"),
+                metrics=per_clip(args.metrics, stem, ".json"),
+                checkpoint=args.checkpoint,
+                force=args.force,
+                cuda_memory_fraction=args.cuda_memory_fraction,
+            )
+        )
+    return cfgs
+
+
+def config_from_args(args: argparse.Namespace) -> RunConfig:
+    """Single-clip form, kept for callers that build a Namespace themselves."""
+    return configs_from_args(args)[0]
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    logging.basicConfig(
+        level=logging.INFO if args.verbose else logging.WARNING, format="%(levelname)s %(name)s: %(message)s"
+    )
+    try:
+        cfgs = configs_from_args(args)
+        if len(cfgs) == 1 and not _is_batch(args):
+            run(cfgs[0], StderrReporter())
+            return 0
+        results = run_many(cfgs, StderrReporter())
+        failures = [r for r in results if isinstance(r, ReshotError)]
+        if failures:
+            print(f"reshot: {len(failures)} of {len(results)} clips failed", file=sys.stderr)
+            return failures[0].exit_code
+        return 0
+    except ReshotError as exc:
+        print(f"reshot: {exc}", file=sys.stderr)
+        return exc.exit_code
+    except KeyboardInterrupt:
+        print("reshot: interrupted", file=sys.stderr)
+        return 130
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main())

reshot/config.py

@@ -0,0 +1,62 @@
+"""Everything a run needs, in one immutable object. Built by the CLI from argv, or by
+library users directly. Validation that does not need I/O happens in `__post_init__`
+so mistakes surface before any work starts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from .errors import InputError
+from .targets import TARGETS
+
+MODEL_VARIANTS = ("small", "base", "large")
+QUALITIES = ("fast", "full")  # model working size: fast = 644×364 for 16:9 (default), full = 924×518
+BACKENDS = ("vda", "fake")
+
+
+@dataclass(frozen=True)
+class RunConfig:
+    input: Path
+    output: Path
+    model: str = "small"
+    backend: str = "vda"
+    device: str = "auto"
+    target: str = "none"
+    fps: float | None = None  # None = preset or source
+    max_res: int = 1280  # cap on the OUTPUT's longer side
+    max_frames: int | None = None
+    quality: str = "fast"  # fast | full — see QUALITIES; `input_size` overrides it
+    input_size: int | None = None  # model short side (multiple of 14); expert override of `quality`
+    invert: bool = False
+    clip_percent: float = 0.0
+    gamma: float = 1.0
+    crf: int = 12
+    npz: Path | None = None
+    metrics: Path | None = None
+    checkpoint: Path | None = None
+    force: bool = False  # ignore the RAM budget
+    cuda_memory_fraction: float | None = None  # cap VRAM to this share of the card (verification aid)
+    extra: dict = field(default_factory=dict)  # free-form, echoed into metrics
+
+    def __post_init__(self) -> None:
+        if self.target not in TARGETS:
+            raise InputError(f"unknown --target {self.target!r}; choose from {', '.join(sorted(TARGETS))}")
+        if self.model not in MODEL_VARIANTS:
+            raise InputError(f"unknown --model {self.model!r}; choose from {', '.join(MODEL_VARIANTS)}")
+        if self.backend not in BACKENDS:
+            raise InputError(f"unknown backend {self.backend!r}; choose from {', '.join(BACKENDS)}")
+        if self.quality not in QUALITIES:
+            raise InputError(f"unknown --quality {self.quality!r}; choose from {', '.join(QUALITIES)}")
+        if self.input_size is not None and (self.input_size < 14 or self.input_size % 14):
+            raise InputError("--input-size must be a positive multiple of 14 (ViT patch size)")
+        if not 0 <= self.clip_percent < 50:
+            raise InputError("--clip must be in [0, 50)")
+        if self.gamma <= 0:
+            raise InputError("--gamma must be > 0")
+        if not 0 <= self.crf <= 51:
+            raise InputError("--crf must be in [0, 51]")
+        if self.max_frames is not None and self.max_frames < 1:
+            raise InputError("--max-frames must be ≥ 1")
+        if self.cuda_memory_fraction is not None and not 0 < self.cuda_memory_fraction <= 1:
+            raise InputError("--cuda-memory-fraction must be in (0, 1]")

reshot/errors.py

@@ -0,0 +1,45 @@
+"""User-facing errors. The CLI prints `str(exc)` and maps each to an exit code; library
+users can catch `ReshotError` for anything the tool considers a *user* problem
+(bad input, over budget, missing tool) as opposed to a bug."""
+
+from __future__ import annotations
+
+
+class ReshotError(Exception):
+    """Base class; `exit_code` is what the CLI returns."""
+
+    exit_code = 1
+
+
+class InputError(ReshotError):
+    """The input video could not be opened or decoded."""
+
+    exit_code = 2
+
+
+class RamBudgetError(ReshotError):
+    """The planned run would exceed the host-RAM budget (see planning.py)."""
+
+    exit_code = 3
+
+    def __init__(self, estimate: int, budget: int, max_frames_ok: int) -> None:
+        self.estimate, self.budget, self.max_frames_ok = estimate, budget, max_frames_ok
+        from .planning import gib
+
+        super().__init__(
+            f"estimated host RAM {gib(estimate)} exceeds the budget {gib(budget)}. "
+            f"Try --max-frames {max_frames_ok} (or split the clip), a smaller --input-size, "
+            "or --force if you know the machine can take it."
+        )
+
+
+class ToolMissingError(ReshotError):
+    """A required external tool (ffmpeg) is missing."""
+
+    exit_code = 4
+
+
+class BackendError(ReshotError):
+    """Unknown backend / model variant, or weights could not be loaded."""
+
+    exit_code = 5