scripts/smoke_gpu.py

"""Real-weight loader regression and GPU smoke test; run from the repo root."""

import sys
from pathlib import Path

import cv2
import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from inference import load_model, predict_mattes

torch.set_num_threads(4)
model = load_model()
capture = cv2.VideoCapture("293085.mp4")
ok, bgr = capture.read()
capture.release()
assert ok, "The supplied example must decode"
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
mattes = list(predict_mattes(model, [rgb, rgb]))
assert mattes[0].shape == rgb.shape[:2]
assert np.isfinite(mattes).all()
assert mattes[0].min() < 0.01 and mattes[0].max() > 0.99
restarted = next(predict_mattes(model, [rgb]))
np.testing.assert_allclose(mattes[0], restarted, atol=1e-3)
print(
    f"PASS: real VRMBG-3.0 weights, CUDA inference and temporal reset on {torch.cuda.get_device_name()}"
)