Spaces:
Running on Zero
Running on Zero
File size: 6,307 Bytes
aa831fd 2ed6676 aa831fd 2ed6676 aa831fd 2ed6676 aa831fd 2ed6676 aa831fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | """Compile FP16 once at startup for the GPU allocated by ZeroGPU."""
import spaces
import logging
import os
from pathlib import Path
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 = None
def __call__(self, *args, **kwargs):
# Cache on the instance, not a ContextVar: ZeroGPU runs each request in
# a fresh copy of the worker's context, which would reload every time.
if self.loaded 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
started = time.perf_counter()
with _register_aoti_cleanup():
self.loaded = torch._inductor.aoti_load_package(self.package_path)
logger.info(
"loaded FP16 package load_s=%.1f", time.perf_counter() - started
)
return self.loaded(*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
|