compilation.py

"""Compile FP16 once at startup for the GPU allocated by ZeroGPU."""

import spaces

import logging
import os
from pathlib import Path
from contextvars import ContextVar
import tempfile
import time

import cv2
import numpy as np
import torch

from inference import predict_mattes

logger = logging.getLogger("vrmbg")


class PackagedFP16:
    """Keep every compiler constant in the package; load once per GPU worker."""

    def __init__(self, package_path):
        self.package_path = package_path
        self.loaded = ContextVar("compiled_fp16", default=None)

    def __call__(self, *args, **kwargs):
        compiled = self.loaded.get()
        if compiled is None:
            # Reuse the pinned Spaces helper's cleanup registration: ZeroGPU
            # kills workers without running PyTorch's temporary-file finalizers.
            from spaces.zero.torch.aoti import _register_aoti_cleanup

            with _register_aoti_cleanup():
                compiled = torch._inductor.aoti_load_package(self.package_path)
            self.loaded.set(compiled)
        return compiled(*args, **kwargs)

    def __reduce__(self):
        # The shared container filesystem holds the package. Never serialize a
        # live CUDA runner back from the startup worker into the parent process.
        return PackagedFP16, (self.package_path,)


def check_mattes(reference, candidate):
    """Reject an invalid compiler result before replacing the eager model."""
    if candidate.shape != reference.shape or not np.isfinite(candidate).all():
        raise RuntimeError("Compiled FP16 returned invalid alpha mattes.")
    error = np.abs(candidate - reference)
    if error.mean() > 0.002 or np.quantile(error, 0.99) > 0.02:
        raise RuntimeError("Compiled FP16 failed the alpha fidelity check.")


def apply_compiled_model(compiled, model):
    """Release eager storage, including integer buffers unsupported by aoti_apply."""
    model.forward = compiled
    for module in model.modules():
        # Replace objects so ZeroGPU's weak CUDA aliases release old storage.
        for name, parameter in module.named_parameters(
            recurse=False, remove_duplicate=False
        ):
            module.register_parameter(
                name,
                torch.nn.Parameter(
                    parameter.new_empty(0), requires_grad=parameter.requires_grad
                ),
            )
        for name, buffer in module.named_buffers(recurse=False, remove_duplicate=False):
            module.register_buffer(name, buffer.new_empty(0))


def compile_fp16(model):
    """Called inside a GPU lease; capture the production input and check output."""
    started = time.perf_counter()
    logger.info("compiling FP16 gpu=%s", torch.cuda.get_device_name())
    cap = cv2.VideoCapture(str(Path(__file__).with_name("293085.mp4")))
    frames = []
    try:
        for _ in range(8):
            ok, bgr = cap.read()
            if not ok:
                raise RuntimeError("Could not read compilation validation frames.")
            frames.append(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
    finally:
        cap.release()

    # Capture through the real preprocessing path: FP16, channels-last,
    # [1, 6, 1024, 1024], including the request-local temporal prior.
    # https://huggingface.co/blog/zerogpu-aoti
    with torch.inference_mode():
        with spaces.aoti_capture(model) as call:
            next(predict_mattes(model, frames[:1]))
        exported = torch.export.export(model, args=call.args, kwargs=call.kwargs)
        cache = Path(__file__).parent / ".cache" / "aoti"
        cache.mkdir(parents=True, exist_ok=True)
        package = Path(tempfile.mkdtemp(prefix="fp16-", dir=cache)) / "model.pt2"
        # Spaces 0.51.3 external weight binding loses this model's generated
        # tensor constants on torch 2.8. Embed them using PyTorch's package API.
        # https://docs.pytorch.org/docs/2.8/torch.compiler_aot_inductor.html
        from spaces.zero.elf import clear_elf_execstack
        from torch._inductor.package.package import package_aoti

        artifacts = torch._inductor.aot_compile(
            exported.module(),
            call.args,
            call.kwargs,
            options={
                "aot_inductor.package": True,
                "aot_inductor.package_constants_in_so": True,
            },
        )
        for artifact in artifacts:
            if artifact.endswith(".so"):
                clear_elf_execstack(Path(artifact))
        package_aoti(str(package), artifacts)
        compiled = PackagedFP16(str(package))
        reference = np.stack(list(predict_mattes(model, frames)))
        candidate = np.stack(list(predict_mattes(compiled, frames)))
        check_mattes(reference, candidate)
    logger.info("compiled FP16 validated compile_s=%.1f", time.perf_counter() - started)
    return compiled


def enable_compiled_model(model):
    """Enable AOT on ZeroGPU; VRMBG_COMPILE=0 restores eager startup."""
    enabled = os.environ.get("VRMBG_COMPILE", os.environ.get("SPACES_ZERO_GPU", "0"))
    if enabled.lower() not in {"1", "true", "t"}:
        logger.info("inference_backend=eager_fp16")
        return False

    # The decorator is a no-op off ZeroGPU, allowing the same path to be tested
    # on a local Linux GPU. Compilation never runs in an inference request.
    @spaces.GPU(duration=600)
    def compile_on_gpu():
        return compile_fp16(model)

    try:
        compiled = compile_on_gpu()
    except Exception:
        # Only a separate ZeroGPU worker can fail without poisoning the parent
        # CUDA context. Explicit local compilation errors must stop that run.
        if os.environ.get("SPACES_ZERO_GPU", "0").lower() not in {"1", "true", "t"}:
            raise
        logger.exception("FP16 compilation failed; using eager FP16")
        return False

    # Apply in the parent process, after the GPU worker has returned its package.
    # The package owns the weights. Do not hide an apply failure after mutation.
    apply_compiled_model(compiled, model)
    logger.info("inference_backend=aot_fp16")
    return True