From 27b5c73bd4c94a551107ffaba374462ef5c0b70d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=9E=8D=E9=9B=AA=E6=97=B6=E5=88=86?= Date: Fri, 8 Dec 2023 18:39:24 +0800 Subject: [PATCH] Fix python binding, add pytorch integration and supports denoising gpu tensors (#46) * fix py binding for newer version of python * add gpu denoise implementation --- CMakeLists.txt | 13 ++++----- README.md | 2 +- common/build/FindPyTorch.cmake | 34 ++++++++++++++++++++++++ common/scripts/README.md | 12 ++++++++- common/scripts/examples/denoise.py | 6 ++--- common/scripts/krr.py | 10 +++++-- common/scripts/run.py | 3 ++- common/scripts/test.py | 23 ++++++++++++++++ src/CMakeLists.txt | 11 +++++--- src/core/config.in.h | 3 +++ src/core/python/common.cpp | 12 +++++++++ src/core/python/py.cpp | 42 +++++++++++++++++++++++++++++- src/ext/CMakeLists.txt | 9 ++++++- src/ext/pybind11 | 2 +- src/main/kiraray.cpp | 2 +- src/render/color.h | 1 + src/util/check.h | 4 ++- 17 files changed, 167 insertions(+), 22 deletions(-) create mode 100644 common/build/FindPyTorch.cmake create mode 100644 common/scripts/test.py create mode 100644 src/core/python/common.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 0fb4362c..62d09c25 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -64,7 +64,6 @@ add_compile_definitions ("$<$:KRR_DEBUG_BUILD>") ############################################################################### enable_language (CUDA) -find_package (CUDA REQUIRED) set(CMAKE_CUDA_STANDARD 17) set(CMAKE_CUDA_STANDARD_REQUIRED ON) @@ -138,13 +137,15 @@ target_compile_options ( # python3 (for building python binding) ############################################################################### list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/src/ext/pybind11/tools") -find_package(PythonLibsNew) -if (PYTHONLIBS_FOUND) +set(KRR_PYTHON_EXECUTABLE "" CACHE STRING "Path to selected python executable") +set(Python_ROOT_DIR ${KRR_PYTHON_EXECUTABLE}) +find_package(Python COMPONENTS Interpreter Development) +if (Python_FOUND) set(KRR_ENABLE_PYTHON ON) - set(PYBIND11_PYTHON_VERSION ${PYTHON_VERSION}) - message("Find python at ${PYTHON_INCLUDE_DIRS}, selected version: ${PYTHON_VERSION}") + set(PYBIND11_PYTHON_VERSION ${Python_VERSION}) + message("Find python at ${Python_INCLUDE_DIRS} : ${Python_LIBRARIES}, selected version: ${Python_VERSION}") else() - message("Do not find python, disable python binding") + message("Did not find python, disable python binding") endif() # Resolve the atlbase.h issue when building with IDEs other than Visual Studio diff --git a/README.md b/README.md index 1ae5b887..1db84b8f 100644 --- a/README.md +++ b/README.md @@ -60,7 +60,7 @@ build/src/kiraray.exe common/configs/example_cbox.json **Camera controlling.** Dragging `LeftMouse` for orbiting, dragging `Scroll` or `Shift+LeftMouse` for panning. `Scroll` for zooming in/out. -**Python binding.** Several simple interfaces are exposed to python scripting via [pybind11](https://github.com/pybind/pybind11), see [scripts](common/scripts) for details. +**Python binding.** Several simple interfaces are exposed to python scripting via [pybind11](https://github.com/pybind/pybind11), including a OptiX denoiser wrapper for denoising NumPy or PyTorch tensors, see [scripts](common/scripts) for details. ### Galleries diff --git a/common/build/FindPyTorch.cmake b/common/build/FindPyTorch.cmake new file mode 100644 index 00000000..40529aac --- /dev/null +++ b/common/build/FindPyTorch.cmake @@ -0,0 +1,34 @@ +if (DEFINED ENV{TORCH_INSTALL_DIR} OR DEFINED TORCH_INSTALL_DIR) +if (DEFINED ENV{TORCH_INSTALL_DIR}) + set( KRR_PYTORCH_ROOT $ENV{TORCH_INSTALL_DIR} CACHE PATH "PyTorch installation directory.") +else() + set( KRR_PYTORCH_ROOT ${TORCH_INSTALL_DIR} CACHE PATH "PyTorch installation directory.") +endif() +message(STATUS "Found PyTorch installation at ${KRR_PYTORCH_ROOT}.") +SET(KRR_ENABLE_PYTORCH ON CACHE INTERNAL "Enable pytorch-interop support.") +else() +message(STATUS "Did not find pytorch. If you want to enable pytorch-interop support, specify its installation location via TORCH_INSTALL_DIR.") +SET(KRR_ENABLE_PYTORCH OFF CACHE INTERNAL "Enable pytorch-interop support.") +endif() + +if (KRR_ENABLE_PYTORCH) +set(TORCH_BIN_DIR + ${KRR_PYTORCH_ROOT}/lib +) +set(TORCH_LIBS + ${TORCH_BIN_DIR}/c10.lib + ${TORCH_BIN_DIR}/c10_cuda.lib + ${TORCH_BIN_DIR}/torch.lib + ${TORCH_BIN_DIR}/torch_cpu.lib + ${TORCH_BIN_DIR}/torch_cuda.lib + ${TORCH_BIN_DIR}/torch_python.lib +) +set(TORCH_INCLUDE_DIRS + ${KRR_PYTORCH_ROOT}/include + ${KRR_PYTORCH_ROOT}/include/torch/csrc/api/include +) +#c10_cuda.lib torch_cpu.lib torch_cuda.lib +add_library(torch_lib INTERFACE) +target_include_directories(torch_lib INTERFACE ${TORCH_INCLUDE_DIRS}) +target_link_libraries(torch_lib INTERFACE ${TORCH_LIBS}) +endif() \ No newline at end of file diff --git a/common/scripts/README.md b/common/scripts/README.md index 33741401..adea8c3f 100644 --- a/common/scripts/README.md +++ b/common/scripts/README.md @@ -1,5 +1,8 @@ # KiRaRay Python Binding +> After Python 3.8, the search paths of DLL dependencies has been reset. Only the system paths, the directory containing the DLL or PYD file are searched for load-time dependencies. Instead, a new function os.add_dll_directory() was added to supply additional search paths. +Necessary DLL paths are exported from `pykrr_common`, see [krr.py](krr.py) for details. + ### Start from script You can start *KiRaRay* from python script with specified configuration. The configuration should be a python dict object. @@ -16,4 +19,11 @@ Kiraray implements a python wrapper for denoising images with optix's built-in a img_denoised = pykrr.denoise(img_noisy, img_normals, img_albedo) ~~~ -This makes it easy to denoise many image files with python scripts. On my RTX3070 Laptop, denoising an image with 1920x1080 takes approximately 1s, while most of the overhead is the memory copy between host and device. It takes about 25ms when acting as a render pass (see [denoise.cpp](../src/render/passes/denoise.cpp)). \ No newline at end of file +This makes it easy to denoise many image files with python scripts. On my RTX3070 Laptop, denoising an image with 1920x1080 takes approximately 1s, while most of the overhead is the memory copy between host and device. It takes about 25ms when acting as a render pass (see [denoise.cpp](../../src/render/passes/denoise/denoise.cpp)). + +#### Denoising PyTorch Tensor +To enable support for PyTorch, you should define the `TORCH_INSTALL_DIR` environment variable to point to the PyTorch installation directory (see [here](../build/FindPyTorch.cmake) for details). The tensor should be on GPU for no CPU-GPU memory copy. + +~~~Python +img_denoised = pykrr.denoise_pytorch_tensor(img_noisy, img_normals, img_albedo) +~~~ \ No newline at end of file diff --git a/common/scripts/examples/denoise.py b/common/scripts/examples/denoise.py index 7976bafe..961e8e11 100644 --- a/common/scripts/examples/denoise.py +++ b/common/scripts/examples/denoise.py @@ -9,13 +9,13 @@ except: raise Exception("Install numpy for manipulating arrays.") from krr import * +import pykrr def denoise(rgb, normals:np.array=None, albedo:np.array=None)->np.array: - import pykrr if normals is not None: - assert (img_albedo.shape == img_rgb.shape) + assert (normals.shape == rgb.shape) if albedo is not None: - assert (img_normals.shape == img_rgb.shape) + assert (albedo.shape == rgb.shape) return pykrr.denoise(rgb.astype(np.float32), normals, albedo) if __name__ == "__main__": diff --git a/common/scripts/krr.py b/common/scripts/krr.py index ea753179..e0abaee0 100644 --- a/common/scripts/krr.py +++ b/common/scripts/krr.py @@ -6,13 +6,19 @@ ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) os.chdir(ROOT_DIR) -print(ROOT_DIR) - sys.path += [os.path.dirname(pyd) for pyd in glob.iglob(os.path.join(ROOT_DIR, "build*", "**/*.pyd"), recursive=True)] sys.path += [os.path.dirname(pyd) for pyd in glob.iglob(os.path.join(ROOT_DIR, "build*", "**/*.so"), recursive=True)] sys.path += [os.path.dirname(pyd) for pyd in glob.iglob(os.path.join(ROOT_DIR, "out*", "**/*.pyd"), recursive=True)] sys.path += [os.path.dirname(pyd) for pyd in glob.iglob(os.path.join(ROOT_DIR, "out*", "**/*.so"), recursive=True)] +import pykrr_common +print("Vulkan root: ", pykrr_common.vulkan_root) +print("PyTorch root: ", pykrr_common.pytorch_root) + +os.add_dll_directory(os.path.join(pykrr_common.vulkan_root, "Bin")) +os.add_dll_directory(os.path.join(pykrr_common.pytorch_root, "lib")) + +import pykrr if __name__ == "__main__": print(sys.path) \ No newline at end of file diff --git a/common/scripts/run.py b/common/scripts/run.py index 6cc61a7d..838399de 100644 --- a/common/scripts/run.py +++ b/common/scripts/run.py @@ -4,6 +4,7 @@ import argparse from krr import * +print(sys.path) try: import numpy as np except: @@ -12,7 +13,7 @@ if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--config", type=str, default="common/configs/example.json", help="Path to scene file") + parser.add_argument("--config", type=str, default="common/configs/example_cbox.json", help="Path to scene file") args = parser.parse_args() config = json.load(open(args.config)) diff --git a/common/scripts/test.py b/common/scripts/test.py new file mode 100644 index 00000000..c05e296f --- /dev/null +++ b/common/scripts/test.py @@ -0,0 +1,23 @@ +import os +import sys +import torch +import argparse + +try: import numpy as np +except: raise Exception("Install numpy for manipulating arrays.") + +from krr import * +import pykrr + +def denoise_tensor(rgb:torch.Tensor, normals:torch.Tensor=None, albedo:torch.Tensor=None)->np.array: + if normals is not None: + assert (rgb.shape == normals.shape) + if albedo is not None: + assert (rgb.shape == albedo.shape) + return pykrr.denoise_torch_tensor(rgb.astype(np.float32), normals, albedo) + +if __name__ == "__main__": + a = torch.rand(size=[720, 1280, 3]) + print(a) + b = denoise_tensor(a) + print(b) \ No newline at end of file diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index ca2f3fad..99529e8d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -8,6 +8,7 @@ MESSAGE ("Source directory: ${KRR_RENDER_SOURCE_DIR}") MESSAGE ("Build output directory: ${CMAKE_BINARY_DIR}") MESSAGE ("CUDA include directory: ${CUDA_INCLUDE_DIRS}") MESSAGE ("Optix include directory: ${OptiX_INCLUDE_DIR}") +MESSAGE ("Vulkan SDK directory: ${KRR_VULKAN_ROOT}") INCLUDE (${KRR_RENDER_ROOT}/common/build/source.cmake) @@ -23,6 +24,7 @@ ADD_DEPENDENCIES (KRR_PTX krr_soa_generated) SET(KRR_SECONDARY_LIBS_ALL ${KRR_SECONDARY_LIBS_ALL} cuda + cudart cublas krr_ext ${CUDA_LIBRARIES} @@ -66,10 +68,13 @@ copy_post_build(krr_lib "$") ADD_SUBDIRECTORY (${KRR_RENDER_SOURCE_DIR}/misc) IF (KRR_ENABLE_PYTHON) - ADD_LIBRARY(pykrr SHARED ${KRR_RENDER_SOURCE_DIR}/core/python/py.cpp) - TARGET_INCLUDE_DIRECTORIES(pykrr SYSTEM PUBLIC ${KRR_INCLUDE_ALL} ${pybind11_INCLUDE_DIR} ${PYTHON_INCLUDE_DIRS}) - TARGET_LINK_LIBRARIES(pykrr PUBLIC ${KRR_LIBRARIES} ${PYTHON_LIBRARIES} pybind11::module -WHOLEARCHIVE:$) + pybind11_add_module(pykrr_common ${KRR_RENDER_SOURCE_DIR}/core/python/common.cpp) + add_library(pykrr MODULE ${KRR_RENDER_SOURCE_DIR}/core/python/py.cpp) + TARGET_LINK_LIBRARIES(pykrr PUBLIC ${KRR_LIBRARIES} -WHOLEARCHIVE:$ + pybind11::module pybind11::lto pybind11::windows_extras) pybind11_extension(pykrr) + set_target_properties(pykrr PROPERTIES CXX_VISIBILITY_PRESET "hidden" + CUDA_VISIBILITY_PRESET "hidden") ENDIF() ADD_EXECUTABLE ( kiraray ${CMAKE_CURRENT_SOURCE_DIR}/main/kiraray.cpp) diff --git a/src/core/config.in.h b/src/core/config.in.h index 6f0a4277..13234e47 100644 --- a/src/core/config.in.h +++ b/src/core/config.in.h @@ -3,9 +3,12 @@ #define KRR_PROJECT_NAME "${CMAKE_PROJECT_NAME}" #define KRR_PROJECT_DIR "${CMAKE_SOURCE_DIR}" #define KRR_BUILD_TYPE "${CMAKE_BUILD_TYPE}" +#define KRR_VULKAN_ROOT "${KRR_VULKAN_ROOT}" +#define KRR_PYTORCH_ROOT "${KRR_PYTORCH_ROOT}" #cmakedefine01 KRR_BUILD_STARLIGHT #cmakedefine01 KRR_RENDER_SPECTRAL +#cmakedefine01 KRR_ENABLE_PYTORCH #define KRR_ENABLE_PROFILE 1 diff --git a/src/core/python/common.cpp b/src/core/python/common.cpp new file mode 100644 index 00000000..fa8008ab --- /dev/null +++ b/src/core/python/common.cpp @@ -0,0 +1,12 @@ +#include "py.h" +#include "common.h" + +KRR_NAMESPACE_BEGIN + +PYBIND11_MODULE(pykrr_common, m) { + // used to find necessary dlls... + m.attr("vulkan_root") = KRR_VULKAN_ROOT; + m.attr("pytorch_root") = KRR_PYTORCH_ROOT; +} + +KRR_NAMESPACE_END \ No newline at end of file diff --git a/src/core/python/py.cpp b/src/core/python/py.cpp index bd5ec8ff..2517686d 100644 --- a/src/core/python/py.cpp +++ b/src/core/python/py.cpp @@ -1,4 +1,8 @@ #include "common.h" +#if KRR_ENABLE_PYTORCH +#define TORCH_API_INCLUDE_EXTENSION_H +#include +#endif #include "main/renderer.h" #include "scene/importer.h" @@ -14,7 +18,7 @@ KRR_NAMESPACE_BEGIN void run(const json& config) { - if (!gpContext) gpContext = std::make_unique(); + if (!gpContext) gpContext = std::make_unique(); { RenderApp app; app.loadConfig(config); @@ -71,6 +75,37 @@ py::array_t denoise(py::array_t normals, + std::optional albedo) { + static bool initialized{}; + static DenoiseBackend denoiser; + if (!gpContext) gpContext = std::make_unique(); + + Vector2i size = {(int) rgb.size(1), (int) rgb.size(0)}; + Log(Info, "Processing image with %lld channels...", rgb.size(2)); + if (rgb.size(2) != 3 && rgb.size(2) != 4) logError("Incorrect image color channels (not 3)!"); + + DenoiseBackend::PixelFormat pixelFormat = rgb.size(2) == 3 + ? DenoiseBackend::PixelFormat::FLOAT3 + : DenoiseBackend::PixelFormat::FLOAT4; + + bool hasGeometry = normals.has_value() && albedo.has_value(); + auto result = torch::empty_like(rgb); + + denoiser.resize(size); + denoiser.setPixelFormat(pixelFormat); + denoiser.setHaveGeometry(hasGeometry); + + denoiser.denoise((CUstream) 0, (float *) rgb.data_ptr(), + hasGeometry ? (float *) normals.value().data_ptr() : nullptr, + hasGeometry ? (float *) albedo.value().data_ptr() : nullptr, + (float *) result.data_ptr()); + return result; +} +#endif + PYBIND11_MODULE(pykrr, m) { m.doc() = "KiRaRay python binding!"; @@ -81,6 +116,11 @@ PYBIND11_MODULE(pykrr, m) { m.def("denoise", &denoise, "Denoise the hdr image using optix's builtin denoiser", "rgb"_a, "normals"_a = py::none(), "albedo"_a = py::none()); +#if KRR_ENABLE_PYTORCH + m.def("denoise_torch_tensor", &denoise_torch_tensor, + "Denoise the hdr image in tensor using optix's builtin denoiser", "rgb"_a, + "normals"_a = py::none(), "albedo"_a = py::none()); +#endif } KRR_NAMESPACE_END \ No newline at end of file diff --git a/src/ext/CMakeLists.txt b/src/ext/CMakeLists.txt index cd999ccc..79cb77eb 100644 --- a/src/ext/CMakeLists.txt +++ b/src/ext/CMakeLists.txt @@ -2,7 +2,10 @@ # GAPI dependencies ############################################################################### INCLUDE (${KRR_RENDER_ROOT}/common/build/FindOptiX.cmake) -find_package(Vulkan REQUIRED) +INCLUDE (${KRR_RENDER_ROOT}/common/build/FindPyTorch.cmake) +FIND_PACKAGE(Vulkan REQUIRED) +CMAKE_PATH(GET Vulkan_INCLUDE_DIR PARENT_PATH VULKAN_ROOT) +SET(KRR_VULKAN_ROOT ${VULKAN_ROOT} CACHE PATH "Path to dected vulkan installation.") ############################################################################### # third party libraries @@ -91,6 +94,10 @@ TARGET_LINK_LIBRARIES(krr_ext INTERFACE openvdb ) +if (KRR_ENABLE_PYTORCH) +TARGET_LINK_LIBRARIES(krr_ext INTERFACE torch_lib) +endif() + # shaderc for runtime compilation of glsl(shaderc) and hlsl(dxc) shaders # add_library(shaderc UNKNOWN IMPORTED) # shaderc is not used currently. # if(WIN32) diff --git a/src/ext/pybind11 b/src/ext/pybind11 index be97c5a9..8b03ffa7 160000 --- a/src/ext/pybind11 +++ b/src/ext/pybind11 @@ -1 +1 @@ -Subproject commit be97c5a98b4b252c524566f508b5c79410d118c6 +Subproject commit 8b03ffa7c06cd9c8a38297b1c8923695d1ff1b07 diff --git a/src/main/kiraray.cpp b/src/main/kiraray.cpp index fd2872f8..41068800 100644 --- a/src/main/kiraray.cpp +++ b/src/main/kiraray.cpp @@ -14,7 +14,7 @@ extern "C" int main(int argc, char *argv[]) { "Switch to Release build for normal performance!"); #endif - string configFile = "common/configs/example_cbox.json"; + string configFile = "common/configs/test/cboxexplosion.json"; if (argc < 2){ Log(Warning, "No config file specified, using default config file: %s", configFile.c_str()); } else { diff --git a/src/render/color.h b/src/render/color.h index 46cf9fcb..7e685156 100644 --- a/src/render/color.h +++ b/src/render/color.h @@ -1,5 +1,6 @@ #pragma once +#include "common.h" #include "krrmath/math.h" #include "util/math_utils.h" diff --git a/src/util/check.h b/src/util/check.h index 01afd77f..c1e0f2de 100644 --- a/src/util/check.h +++ b/src/util/check.h @@ -58,7 +58,9 @@ KRR_NAMESPACE_BEGIN } \ } while (0) - +#ifdef CHECK +#undef CHECK +#endif #define CHECK(x) assert(x) #define CHECK_IMPL(a, b, op) assert((a)op(b))