File size: 1,561 Bytes
de4d6ae 1f9a369 de4d6ae 1f9a369 de4d6ae | 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 | """Load pathtracer-diff natively on this machine (Windows/SAURON included).
The published Hugging Face kernel carries `*-linux` variants only, so on
Windows `get_kernel` has no matching build. The source is right here, though:
JIT-build it with `torch.utils.cpp_extension` and expose the identical
`Scene` / `Camera` / `render` API.
import load_local
ptd = load_local.load()
"""
import os
import sys
import types
_cached = None
def load(verbose=False):
"""Return the pathtracer_diff module, JIT-built from the local source."""
global _cached
if _cached is not None:
return _cached
from torch.utils.cpp_extension import load as _jit
root = os.path.dirname(os.path.abspath(__file__))
cuda_dir = os.path.join(root, "pathtracer_diff_cuda")
ext = _jit(name="pathtracer_diff_ext",
sources=[os.path.join(root, "local_bind.cpp"),
os.path.join(cuda_dir, "pathtracer.cu"),
os.path.join(cuda_dir, "geometry_grad.cu")],
extra_include_paths=[cuda_dir],
verbose=verbose)
ops_mod = types.ModuleType("pathtracer_diff._ops")
ops_mod.ops = type("_Ops", (), {
"pt_forward": staticmethod(ext.pt_forward),
"pt_backward": staticmethod(ext.pt_backward),
"pt_geometry_grad": staticmethod(ext.pt_geometry_grad),
})()
sys.modules["pathtracer_diff._ops"] = ops_mod
sys.path.insert(0, os.path.join(root, "torch-ext"))
import pathtracer_diff
_cached = pathtracer_diff
return pathtracer_diff
|