pathtracer-diff / load_local.py
phanerozoic's picture
v3: geometry gradients (dual-number interior + shadow and camera silhouette edge sampling)
1f9a369 verified
Raw
History Blame
1.56 kB
"""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