# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. # # NVIDIA CORPORATION and its licensors retain all intellectual property # and proprietary rights in and to this software, related documentation # and any modifications thereto. Any use, reproduction, disclosure or # distribution of this software and related documentation without an express # license agreement from NVIDIA CORPORATION is strictly prohibited. import importlib import logging import numpy as np import os import torch import torch.utils.cpp_extension from . import _C #---------------------------------------------------------------------------- # C++/Cuda plugin compiler/loader. _cached_plugin = {} def _get_plugin(gl=False): assert isinstance(gl, bool) # Modified with precompiled torch CUDA extension if not gl: return _C # Return cached plugin if already loaded. if _cached_plugin.get(gl, None) is not None: return _cached_plugin[gl] # Make sure we can find the necessary compiler and libary binaries. if os.name == 'nt': lib_dir = os.path.dirname(__file__) + r"\..\lib" def find_cl_path(): import glob def get_sort_key(x): # Primary criterion is VS version, secondary is edition, third is internal MSVC version. x = x.split('\\')[3:] x[1] = {'BuildTools': '~0', 'Community': '~1', 'Pro': '~2', 'Professional': '~3', 'Enterprise': '~4'}.get(x[1], x[1]) return x vs_relative_path = r"\Microsoft Visual Studio\*\*\VC\Tools\MSVC\*\bin\Hostx64\x64" paths = glob.glob(r"C:\Program Files" + vs_relative_path) paths += glob.glob(r"C:\Program Files (x86)" + vs_relative_path) if paths: return sorted(paths, key=get_sort_key)[-1] # If cl.exe is not on path, try to find it. if os.system("where cl.exe >nul 2>nul") != 0: cl_path = find_cl_path() if cl_path is None: raise RuntimeError("Could not locate a supported Microsoft Visual C++ installation") os.environ['PATH'] += ';' + cl_path # Compiler options. common_opts = ['-DNVDR_TORCH'] cc_opts = [] if os.name == 'nt': cc_opts += ['/wd4067', '/wd4624'] # Disable warnings in torch headers. # Linker options for the GL-interfacing plugin. ldflags = [] if gl: if os.name == 'posix': ldflags = ['-lGL', '-lEGL'] elif os.name == 'nt': libs = ['gdi32', 'opengl32', 'user32', 'setgpu'] ldflags = ['/LIBPATH:' + lib_dir] + ['/DEFAULTLIB:' + x for x in libs] # List of source files. if gl: source_files = [ '../common/common.cpp', '../common/glutil.cpp', '../common/rasterize_gl.cpp', 'torch_bindings_gl.cpp', 'torch_rasterize_gl.cpp', ] else: source_files = [ '../common/cudaraster/impl/Buffer.cpp', '../common/cudaraster/impl/CudaRaster.cpp', '../common/cudaraster/impl/RasterImpl.cu', '../common/cudaraster/impl/RasterImpl.cpp', '../common/common.cpp', '../common/rasterize.cu', '../common/interpolate.cu', '../common/texture.cu', '../common/texture.cpp', '../common/antialias.cu', 'torch_bindings.cpp', 'torch_rasterize.cpp', 'torch_interpolate.cpp', 'torch_texture.cpp', 'torch_antialias.cpp', ] # Some containers set this to contain old architectures that won't compile. We only need the one installed in the machine. os.environ['TORCH_CUDA_ARCH_LIST'] = '' # On Linux, show a warning if GLEW is being forcibly loaded when compiling the GL plugin. if gl and (os.name == 'posix') and ('libGLEW' in os.environ.get('LD_PRELOAD', '')): logging.getLogger('nvdiffrast').warning("Warning: libGLEW is being loaded via LD_PRELOAD, and will probably conflict with the OpenGL plugin") # Try to detect if a stray lock file is left in cache directory and show a warning. This sometimes happens on Windows if the build is interrupted at just the right moment. plugin_name = 'nvdiffrast_plugin' + ('_gl' if gl else '') try: lock_fn = os.path.join(torch.utils.cpp_extension._get_build_directory(plugin_name, False), 'lock') if os.path.exists(lock_fn): logging.getLogger('nvdiffrast').warning("Lock file exists in build directory: '%s'" % lock_fn) except: pass # Speed up compilation on Windows. if os.name == 'nt': # Skip telemetry sending step in vcvarsall.bat os.environ['VSCMD_SKIP_SENDTELEMETRY'] = '1' # Opportunistically patch distutils to cache MSVC environments. try: import distutils._msvccompiler import functools if not hasattr(distutils._msvccompiler._get_vc_env, '__wrapped__'): distutils._msvccompiler._get_vc_env = functools.lru_cache()(distutils._msvccompiler._get_vc_env) except: pass # Compile and load. source_paths = [os.path.join(os.path.dirname(__file__), fn) for fn in source_files] torch.utils.cpp_extension.load(name=plugin_name, sources=source_paths, extra_cflags=common_opts+cc_opts, extra_cuda_cflags=common_opts+['-lineinfo'], extra_ldflags=ldflags, with_cuda=True, verbose=False) # Import, cache, and return the compiled module. _cached_plugin[gl] = importlib.import_module(plugin_name) return _cached_plugin[gl] #---------------------------------------------------------------------------- # Log level. #---------------------------------------------------------------------------- def get_log_level(): '''Get current log level. Returns: Current log level in nvdiffrast. See `set_log_level()` for possible values. ''' return _get_plugin().get_log_level() def set_log_level(level): '''Set log level. Log levels follow the convention on the C++ side of Torch: 0 = Info, 1 = Warning, 2 = Error, 3 = Fatal. The default log level is 1. Args: level: New log level as integer. Internal nvdiffrast messages of this severity or higher will be printed, while messages of lower severity will be silent. ''' _get_plugin().set_log_level(level) #---------------------------------------------------------------------------- # CudaRaster state wrapper. #---------------------------------------------------------------------------- class RasterizeCudaContext: def __init__(self, device=None): '''Create a new Cuda rasterizer context. The context is deleted and internal storage is released when the object is destroyed. Args: device (Optional): Cuda device on which the context is created. Type can be `torch.device`, string (e.g., `'cuda:1'`), or int. If not specified, context will be created on currently active Cuda device. Returns: The newly created Cuda rasterizer context. ''' if device is None: cuda_device_idx = torch.cuda.current_device() else: with torch.cuda.device(device): cuda_device_idx = torch.cuda.current_device() self.cpp_wrapper = _get_plugin().RasterizeCRStateWrapper(cuda_device_idx) self.output_db = True self.active_depth_peeler = None #---------------------------------------------------------------------------- # GL state wrapper. #---------------------------------------------------------------------------- class RasterizeGLContext: def __init__(self, output_db=True, mode='automatic', device=None): '''Create a new OpenGL rasterizer context. Creating an OpenGL context is a slow operation so you should usually reuse the same context in all calls to `rasterize()` on the same CPU thread. The OpenGL context is deleted when the object is destroyed. Side note: When using the OpenGL context in a rasterization operation, the context's internal framebuffer object is automatically enlarged to accommodate the rasterization operation's output shape, but it is never shrunk in size until the context is destroyed. Thus, if you need to rasterize, say, deep low-resolution tensors and also shallow high-resolution tensors, you can conserve GPU memory by creating two separate OpenGL contexts for these tasks. In this scenario, using the same OpenGL context for both tasks would end up reserving GPU memory for a deep, high-resolution output tensor. Args: output_db (bool): Compute and output image-space derivates of barycentrics. mode: OpenGL context handling mode. Valid values are 'manual' and 'automatic'. device (Optional): Cuda device on which the context is created. Type can be `torch.device`, string (e.g., `'cuda:1'`), or int. If not specified, context will be created on currently active Cuda device. Returns: The newly created OpenGL rasterizer context. ''' assert output_db is True or output_db is False assert mode in ['automatic', 'manual'] self.output_db = output_db self.mode = mode if device is None: cuda_device_idx = torch.cuda.current_device() else: with torch.cuda.device(device): cuda_device_idx = torch.cuda.current_device() self.cpp_wrapper = _get_plugin(gl=True).RasterizeGLStateWrapper(output_db, mode == 'automatic', cuda_device_idx) self.active_depth_peeler = None # For error checking only. def set_context(self): '''Set (activate) OpenGL context in the current CPU thread. Only available if context was created in manual mode. ''' assert self.mode == 'manual' self.cpp_wrapper.set_context() def release_context(self): '''Release (deactivate) currently active OpenGL context. Only available if context was created in manual mode. ''' assert self.mode == 'manual' self.cpp_wrapper.release_context() #---------------------------------------------------------------------------- # Rasterize. #---------------------------------------------------------------------------- class _rasterize_func(torch.autograd.Function): @staticmethod def forward(ctx, raster_ctx, pos, tri, resolution, ranges, grad_db, peeling_idx): if isinstance(raster_ctx, RasterizeGLContext): out, out_db = _get_plugin(gl=True).rasterize_fwd_gl(raster_ctx.cpp_wrapper, pos, tri, resolution, ranges, peeling_idx) else: out, out_db = _get_plugin().rasterize_fwd_cuda(raster_ctx.cpp_wrapper, pos, tri, resolution, ranges, peeling_idx) ctx.save_for_backward(pos, tri, out) ctx.saved_grad_db = grad_db return out, out_db @staticmethod def backward(ctx, dy, ddb): pos, tri, out = ctx.saved_tensors if ctx.saved_grad_db: g_pos = _get_plugin().rasterize_grad_db(pos, tri, out, dy, ddb) else: g_pos = _get_plugin().rasterize_grad(pos, tri, out, dy) return None, g_pos, None, None, None, None, None # Op wrapper. def rasterize(glctx, pos, tri, resolution, ranges=None, grad_db=True): '''Rasterize triangles. All input tensors must be contiguous and reside in GPU memory except for the `ranges` tensor that, if specified, has to reside in CPU memory. The output tensors will be contiguous and reside in GPU memory. Args: glctx: Rasterizer context of type `RasterizeGLContext` or `RasterizeCudaContext`. pos: Vertex position tensor with dtype `torch.float32`. To enable range mode, this tensor should have a 2D shape [num_vertices, 4]. To enable instanced mode, use a 3D shape [minibatch_size, num_vertices, 4]. tri: Triangle tensor with shape [num_triangles, 3] and dtype `torch.int32`. resolution: Output resolution as integer tuple (height, width). ranges: In range mode, tensor with shape [minibatch_size, 2] and dtype `torch.int32`, specifying start indices and counts into `tri`. Ignored in instanced mode. grad_db: Propagate gradients of image-space derivatives of barycentrics into `pos` in backward pass. Ignored if using an OpenGL context that was not configured to output image-space derivatives. Returns: A tuple of two tensors. The first output tensor has shape [minibatch_size, height, width, 4] and contains the main rasterizer output in order (u, v, z/w, triangle_id). If the OpenGL context was configured to output image-space derivatives of barycentrics, the second output tensor will also have shape [minibatch_size, height, width, 4] and contain said derivatives in order (du/dX, du/dY, dv/dX, dv/dY). Otherwise it will be an empty tensor with shape [minibatch_size, height, width, 0]. ''' assert isinstance(glctx, (RasterizeGLContext, RasterizeCudaContext)) assert grad_db is True or grad_db is False grad_db = grad_db and glctx.output_db # Sanitize inputs. assert isinstance(pos, torch.Tensor) and isinstance(tri, torch.Tensor) resolution = tuple(resolution) if ranges is None: ranges = torch.empty(size=(0, 2), dtype=torch.int32, device='cpu') else: assert isinstance(ranges, torch.Tensor) # Check that context is not currently reserved for depth peeling. if glctx.active_depth_peeler is not None: return RuntimeError("Cannot call rasterize() during depth peeling operation, use rasterize_next_layer() instead") # Instantiate the function. return _rasterize_func.apply(glctx, pos, tri, resolution, ranges, grad_db, -1) #---------------------------------------------------------------------------- # Depth peeler context manager for rasterizing multiple depth layers. #---------------------------------------------------------------------------- class DepthPeeler: def __init__(self, glctx, pos, tri, resolution, ranges=None, grad_db=True): '''Create a depth peeler object for rasterizing multiple depth layers. Arguments are the same as in `rasterize()`. Returns: The newly created depth peeler. ''' assert isinstance(glctx, (RasterizeGLContext, RasterizeCudaContext)) assert grad_db is True or grad_db is False grad_db = grad_db and glctx.output_db # Sanitize inputs as usual. assert isinstance(pos, torch.Tensor) and isinstance(tri, torch.Tensor) resolution = tuple(resolution) if ranges is None: ranges = torch.empty(size=(0, 2), dtype=torch.int32, device='cpu') else: assert isinstance(ranges, torch.Tensor) # Store all the parameters. self.raster_ctx = glctx self.pos = pos self.tri = tri self.resolution = resolution self.ranges = ranges self.grad_db = grad_db self.peeling_idx = None def __enter__(self): if self.raster_ctx is None: raise RuntimeError("Cannot re-enter a terminated depth peeling operation") if self.raster_ctx.active_depth_peeler is not None: raise RuntimeError("Cannot have multiple depth peelers active simultaneously in a rasterization context") self.raster_ctx.active_depth_peeler = self self.peeling_idx = 0 return self def __exit__(self, *args): assert self.raster_ctx.active_depth_peeler is self self.raster_ctx.active_depth_peeler = None self.raster_ctx = None # Remove all references to input tensor so they're not left dangling. self.pos = None self.tri = None self.resolution = None self.ranges = None self.grad_db = None self.peeling_idx = None return None def rasterize_next_layer(self): '''Rasterize next depth layer. Operation is equivalent to `rasterize()` except that previously reported surface points are culled away. Returns: A tuple of two tensors as in `rasterize()`. ''' assert self.raster_ctx.active_depth_peeler is self assert self.peeling_idx >= 0 result = _rasterize_func.apply(self.raster_ctx, self.pos, self.tri, self.resolution, self.ranges, self.grad_db, self.peeling_idx) self.peeling_idx += 1 return result #---------------------------------------------------------------------------- # Interpolate. #---------------------------------------------------------------------------- # Output pixel differentials for at least some attributes. class _interpolate_func_da(torch.autograd.Function): @staticmethod def forward(ctx, attr, rast, tri, rast_db, diff_attrs_all, diff_attrs_list): out, out_da = _get_plugin().interpolate_fwd_da(attr, rast, tri, rast_db, diff_attrs_all, diff_attrs_list) ctx.save_for_backward(attr, rast, tri, rast_db) ctx.saved_misc = diff_attrs_all, diff_attrs_list return out, out_da @staticmethod def backward(ctx, dy, dda): attr, rast, tri, rast_db = ctx.saved_tensors diff_attrs_all, diff_attrs_list = ctx.saved_misc g_attr, g_rast, g_rast_db = _get_plugin().interpolate_grad_da(attr, rast, tri, dy, rast_db, dda, diff_attrs_all, diff_attrs_list) return g_attr, g_rast, None, g_rast_db, None, None # No pixel differential for any attribute. class _interpolate_func(torch.autograd.Function): @staticmethod def forward(ctx, attr, rast, tri): out, out_da = _get_plugin().interpolate_fwd(attr, rast, tri) ctx.save_for_backward(attr, rast, tri) return out, out_da @staticmethod def backward(ctx, dy, _): attr, rast, tri = ctx.saved_tensors g_attr, g_rast = _get_plugin().interpolate_grad(attr, rast, tri, dy) return g_attr, g_rast, None # Op wrapper. def interpolate(attr, rast, tri, rast_db=None, diff_attrs=None): """Interpolate vertex attributes. All input tensors must be contiguous and reside in GPU memory. The output tensors will be contiguous and reside in GPU memory. Args: attr: Attribute tensor with dtype `torch.float32`. Shape is [num_vertices, num_attributes] in range mode, or [minibatch_size, num_vertices, num_attributes] in instanced mode. Broadcasting is supported along the minibatch axis. rast: Main output tensor from `rasterize()`. tri: Triangle tensor with shape [num_triangles, 3] and dtype `torch.int32`. rast_db: (Optional) Tensor containing image-space derivatives of barycentrics, i.e., the second output tensor from `rasterize()`. Enables computing image-space derivatives of attributes. diff_attrs: (Optional) List of attribute indices for which image-space derivatives are to be computed. Special value 'all' is equivalent to list [0, 1, ..., num_attributes - 1]. Returns: A tuple of two tensors. The first output tensor contains interpolated attributes and has shape [minibatch_size, height, width, num_attributes]. If `rast_db` and `diff_attrs` were specified, the second output tensor contains the image-space derivatives of the selected attributes and has shape [minibatch_size, height, width, 2 * len(diff_attrs)]. The derivatives of the first selected attribute A will be on channels 0 and 1 as (dA/dX, dA/dY), etc. Otherwise, the second output tensor will be an empty tensor with shape [minibatch_size, height, width, 0]. """ # Sanitize the list of pixel differential attributes. if diff_attrs is None: diff_attrs = [] elif diff_attrs != 'all': diff_attrs = np.asarray(diff_attrs, np.int32) assert len(diff_attrs.shape) == 1 diff_attrs = diff_attrs.tolist() diff_attrs_all = int(diff_attrs == 'all') diff_attrs_list = [] if diff_attrs_all else diff_attrs # Check inputs. assert all(isinstance(x, torch.Tensor) for x in (attr, rast, tri)) if diff_attrs: assert isinstance(rast_db, torch.Tensor) # Choose stub. if diff_attrs: return _interpolate_func_da.apply(attr, rast, tri, rast_db, diff_attrs_all, diff_attrs_list) else: return _interpolate_func.apply(attr, rast, tri) #---------------------------------------------------------------------------- # Texture #---------------------------------------------------------------------------- # Linear-mipmap-linear and linear-mipmap-nearest: Mipmaps enabled. class _texture_func_mip(torch.autograd.Function): @staticmethod def forward(ctx, filter_mode, tex, uv, uv_da, mip_level_bias, mip_wrapper, filter_mode_enum, boundary_mode_enum, *mip_stack): empty = torch.tensor([]) if uv_da is None: uv_da = empty if mip_level_bias is None: mip_level_bias = empty if mip_wrapper is None: mip_wrapper = _get_plugin().TextureMipWrapper() out = _get_plugin().texture_fwd_mip(tex, uv, uv_da, mip_level_bias, mip_wrapper, mip_stack, filter_mode_enum, boundary_mode_enum) ctx.save_for_backward(tex, uv, uv_da, mip_level_bias, *mip_stack) ctx.saved_misc = filter_mode, mip_wrapper, filter_mode_enum, boundary_mode_enum return out @staticmethod def backward(ctx, dy): tex, uv, uv_da, mip_level_bias, *mip_stack = ctx.saved_tensors filter_mode, mip_wrapper, filter_mode_enum, boundary_mode_enum = ctx.saved_misc if filter_mode == 'linear-mipmap-linear': g_tex, g_uv, g_uv_da, g_mip_level_bias, g_mip_stack = _get_plugin().texture_grad_linear_mipmap_linear(tex, uv, dy, uv_da, mip_level_bias, mip_wrapper, mip_stack, filter_mode_enum, boundary_mode_enum) return (None, g_tex, g_uv, g_uv_da, g_mip_level_bias, None, None, None) + tuple(g_mip_stack) else: # linear-mipmap-nearest g_tex, g_uv, g_mip_stack = _get_plugin().texture_grad_linear_mipmap_nearest(tex, uv, dy, uv_da, mip_level_bias, mip_wrapper, mip_stack, filter_mode_enum, boundary_mode_enum) return (None, g_tex, g_uv, None, None, None, None, None) + tuple(g_mip_stack) # Linear and nearest: Mipmaps disabled. class _texture_func(torch.autograd.Function): @staticmethod def forward(ctx, filter_mode, tex, uv, filter_mode_enum, boundary_mode_enum): out = _get_plugin().texture_fwd(tex, uv, filter_mode_enum, boundary_mode_enum) ctx.save_for_backward(tex, uv) ctx.saved_misc = filter_mode, filter_mode_enum, boundary_mode_enum return out @staticmethod def backward(ctx, dy): tex, uv = ctx.saved_tensors filter_mode, filter_mode_enum, boundary_mode_enum = ctx.saved_misc if filter_mode == 'linear': g_tex, g_uv = _get_plugin().texture_grad_linear(tex, uv, dy, filter_mode_enum, boundary_mode_enum) return None, g_tex, g_uv, None, None else: # nearest g_tex = _get_plugin().texture_grad_nearest(tex, uv, dy, filter_mode_enum, boundary_mode_enum) return None, g_tex, None, None, None # Op wrapper. def texture(tex, uv, uv_da=None, mip_level_bias=None, mip=None, filter_mode='auto', boundary_mode='wrap', max_mip_level=None): """Perform texture sampling. All input tensors must be contiguous and reside in GPU memory. The output tensor will be contiguous and reside in GPU memory. Args: tex: Texture tensor with dtype `torch.float32`. For 2D textures, must have shape [minibatch_size, tex_height, tex_width, tex_channels]. For cube map textures, must have shape [minibatch_size, 6, tex_height, tex_width, tex_channels] where tex_width and tex_height are equal. Note that `boundary_mode` must also be set to 'cube' to enable cube map mode. Broadcasting is supported along the minibatch axis. uv: Tensor containing per-pixel texture coordinates. When sampling a 2D texture, must have shape [minibatch_size, height, width, 2]. When sampling a cube map texture, must have shape [minibatch_size, height, width, 3]. uv_da: (Optional) Tensor containing image-space derivatives of texture coordinates. Must have same shape as `uv` except for the last dimension that is to be twice as long. mip_level_bias: (Optional) Per-pixel bias for mip level selection. If `uv_da` is omitted, determines mip level directly. Must have shape [minibatch_size, height, width]. mip: (Optional) Preconstructed mipmap stack from a `texture_construct_mip()` call, or a list of tensors specifying a custom mipmap stack. When specifying a custom mipmap stack, the tensors in the list must follow the same format as `tex` except for width and height that must follow the usual rules for mipmap sizes. The base level texture is still supplied in `tex` and must not be included in the list. Gradients of a custom mipmap stack are not automatically propagated to base texture but the mipmap tensors will receive gradients of their own. If a mipmap stack is not specified but the chosen filter mode requires it, the mipmap stack is constructed internally and discarded afterwards. filter_mode: Texture filtering mode to be used. Valid values are 'auto', 'nearest', 'linear', 'linear-mipmap-nearest', and 'linear-mipmap-linear'. Mode 'auto' selects 'linear' if neither `uv_da` or `mip_level_bias` is specified, and 'linear-mipmap-linear' when at least one of them is specified, these being the highest-quality modes possible depending on the availability of the image-space derivatives of the texture coordinates or direct mip level information. boundary_mode: Valid values are 'wrap', 'clamp', 'zero', and 'cube'. If `tex` defines a cube map, this must be set to 'cube'. The default mode 'wrap' takes fractional part of texture coordinates. Mode 'clamp' clamps texture coordinates to the centers of the boundary texels. Mode 'zero' virtually extends the texture with all-zero values in all directions. max_mip_level: If specified, limits the number of mipmaps constructed and used in mipmap-based filter modes. Returns: A tensor containing the results of the texture sampling with shape [minibatch_size, height, width, tex_channels]. Cube map fetches with invalid uv coordinates (e.g., zero vectors) output all zeros and do not propagate gradients. """ # Default filter mode. if filter_mode == 'auto': filter_mode = 'linear-mipmap-linear' if (uv_da is not None or mip_level_bias is not None) else 'linear' # Sanitize inputs. if max_mip_level is None: max_mip_level = -1 else: max_mip_level = int(max_mip_level) assert max_mip_level >= 0 # Check inputs. assert isinstance(tex, torch.Tensor) and isinstance(uv, torch.Tensor) if 'mipmap' in filter_mode: assert isinstance(uv_da, torch.Tensor) or isinstance(mip_level_bias, torch.Tensor) # If mipping disabled via max level=0, we may as well use simpler filtering internally. if max_mip_level == 0 and filter_mode in ['linear-mipmap-nearest', 'linear-mipmap-linear']: filter_mode = 'linear' # Convert filter mode to internal enumeration. filter_mode_dict = {'nearest': 0, 'linear': 1, 'linear-mipmap-nearest': 2, 'linear-mipmap-linear': 3} filter_mode_enum = filter_mode_dict[filter_mode] # Convert boundary mode to internal enumeration. boundary_mode_dict = {'cube': 0, 'wrap': 1, 'clamp': 2, 'zero': 3} boundary_mode_enum = boundary_mode_dict[boundary_mode] # Construct a mipmap if necessary. if 'mipmap' in filter_mode: mip_wrapper, mip_stack = None, [] if mip is not None: assert isinstance(mip, (_get_plugin().TextureMipWrapper, list)) if isinstance(mip, list): assert all(isinstance(x, torch.Tensor) for x in mip) mip_stack = mip else: mip_wrapper = mip else: mip_wrapper = _get_plugin().texture_construct_mip(tex, max_mip_level, boundary_mode == 'cube') # Choose stub. if filter_mode == 'linear-mipmap-linear' or filter_mode == 'linear-mipmap-nearest': return _texture_func_mip.apply(filter_mode, tex, uv, uv_da, mip_level_bias, mip_wrapper, filter_mode_enum, boundary_mode_enum, *mip_stack) else: return _texture_func.apply(filter_mode, tex, uv, filter_mode_enum, boundary_mode_enum) # Mipmap precalculation for cases where the texture stays constant. def texture_construct_mip(tex, max_mip_level=None, cube_mode=False): """Construct a mipmap stack for a texture. This function can be used for constructing a mipmap stack for a texture that is known to remain constant. This avoids reconstructing it every time `texture()` is called. Args: tex: Texture tensor with the same constraints as in `texture()`. max_mip_level: If specified, limits the number of mipmaps constructed. cube_mode: Must be set to True if `tex` specifies a cube map texture. Returns: An opaque object containing the mipmap stack. This can be supplied in a call to `texture()` in the `mip` argument. """ assert isinstance(tex, torch.Tensor) assert cube_mode is True or cube_mode is False if max_mip_level is None: max_mip_level = -1 else: max_mip_level = int(max_mip_level) assert max_mip_level >= 0 return _get_plugin().texture_construct_mip(tex, max_mip_level, cube_mode) #---------------------------------------------------------------------------- # Antialias. #---------------------------------------------------------------------------- class _antialias_func(torch.autograd.Function): @staticmethod def forward(ctx, color, rast, pos, tri, topology_hash, pos_gradient_boost): out, work_buffer = _get_plugin().antialias_fwd(color, rast, pos, tri, topology_hash) ctx.save_for_backward(color, rast, pos, tri) ctx.saved_misc = pos_gradient_boost, work_buffer return out @staticmethod def backward(ctx, dy): color, rast, pos, tri = ctx.saved_tensors pos_gradient_boost, work_buffer = ctx.saved_misc g_color, g_pos = _get_plugin().antialias_grad(color, rast, pos, tri, dy, work_buffer) if pos_gradient_boost != 1.0: g_pos = g_pos * pos_gradient_boost return g_color, None, g_pos, None, None, None # Op wrapper. def antialias(color, rast, pos, tri, topology_hash=None, pos_gradient_boost=1.0): """Perform antialiasing. All input tensors must be contiguous and reside in GPU memory. The output tensor will be contiguous and reside in GPU memory. Note that silhouette edge determination is based on vertex indices in the triangle tensor. For it to work properly, a vertex belonging to multiple triangles must be referred to using the same vertex index in each triangle. Otherwise, nvdiffrast will always classify the adjacent edges as silhouette edges, which leads to bad performance and potentially incorrect gradients. If you are unsure whether your data is good, check which pixels are modified by the antialias operation and compare to the example in the documentation. Args: color: Input image to antialias with shape [minibatch_size, height, width, num_channels]. rast: Main output tensor from `rasterize()`. pos: Vertex position tensor used in the rasterization operation. tri: Triangle tensor used in the rasterization operation. topology_hash: (Optional) Preconstructed topology hash for the triangle tensor. If not specified, the topology hash is constructed internally and discarded afterwards. pos_gradient_boost: (Optional) Multiplier for gradients propagated to `pos`. Returns: A tensor containing the antialiased image with the same shape as `color` input tensor. """ # Check inputs. assert all(isinstance(x, torch.Tensor) for x in (color, rast, pos, tri)) # Construct topology hash unless provided by user. if topology_hash is not None: assert isinstance(topology_hash, _get_plugin().TopologyHashWrapper) else: topology_hash = _get_plugin().antialias_construct_topology_hash(tri) # Instantiate the function. return _antialias_func.apply(color, rast, pos, tri, topology_hash, pos_gradient_boost) # Topology hash precalculation for cases where the triangle array stays constant. def antialias_construct_topology_hash(tri): """Construct a topology hash for a triangle tensor. This function can be used for constructing a topology hash for a triangle tensor that is known to remain constant. This avoids reconstructing it every time `antialias()` is called. Args: tri: Triangle tensor with shape [num_triangles, 3]. Must be contiguous and reside in GPU memory. Returns: An opaque object containing the topology hash. This can be supplied in a call to `antialias()` in the `topology_hash` argument. """ assert isinstance(tri, torch.Tensor) return _get_plugin().antialias_construct_topology_hash(tri) #----------------------------------------------------------------------------