Fix python binding, add pytorch integration and supports denoising gp…

…u tensors (#46)

* fix py binding for newer version of python
* add gpu denoise implementation

CMakeLists.txt

@@ -64,7 +64,6 @@ add_compile_definitions ("$<$<CONFIG:Debug>: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

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
 

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()

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)).
+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)
+~~~

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__":

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)

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))

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)

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 "$<TARGET_FILE:IlmBase::Half>")
 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:$<TARGET_FILE:krr_lib>)
+	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:$<TARGET_FILE:krr_lib>
+		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)

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
 

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

src/core/python/py.cpp

@@ -1,4 +1,8 @@
 #include "common.h"
+#if KRR_ENABLE_PYTORCH
+#define TORCH_API_INCLUDE_EXTENSION_H
+#include <torch/torch.h>
+#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<Context>();
+	if (!gpContext) gpContext = std::make_unique<Context>(); 
 	{
 		RenderApp app;
 		app.loadConfig(config);
@@ -71,6 +75,37 @@ py::array_t<float> denoise(py::array_t<float, py::array::c_style | py::array::fo
 	return result;
 }
 
+#if KRR_ENABLE_PYTORCH
+torch::Tensor denoise_torch_tensor(torch::Tensor rgb, 
+	std::optional<torch::Tensor> normals, 
+	std::optional<torch::Tensor> albedo) {
+	static bool initialized{};
+	static DenoiseBackend denoiser;
+	if (!gpContext) gpContext = std::make_unique<Context>();
+
+	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

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)

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 {

src/render/color.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "common.h"
 #include "krrmath/math.h"
 #include "util/math_utils.h"
 

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))