├── .gitignore ├── .readthedocs.yaml ├── LICENSE ├── README.md ├── diffrp ├── __init__.py ├── __main__.py ├── loaders │ ├── __init__.py │ ├── _trimesh_gltf.py │ └── gltf_loader.py ├── materials │ ├── __init__.py │ ├── base_material.py │ ├── default_material.py │ └── gltf_material.py ├── plugins │ ├── .gitignore │ ├── __init__.py │ └── mikktspace │ │ ├── __init__.py │ │ ├── mikktinterface.c │ │ ├── mikktspace.c │ │ └── mikktspace.h ├── rendering │ ├── __init__.py │ ├── camera.py │ ├── denoiser.py │ ├── interpolator.py │ ├── mixin.py │ ├── path_tracing.py │ └── surface_deferred.py ├── resources │ ├── __init__.py │ ├── denoisers │ │ └── rt_hdr_alb_nrm.pt │ ├── hdri_exrs │ │ └── newport_loft.exr │ ├── hdris.py │ └── luts │ │ └── agx-base-contrast.pt ├── scene │ ├── __init__.py │ ├── lights.py │ ├── objects.py │ └── scene.py ├── scripts │ ├── __init__.py │ ├── main.py │ └── quick_render.py ├── utils │ ├── __init__.py │ ├── cache.py │ ├── colors.py │ ├── composite.py │ ├── coordinates.py │ ├── exchange.py │ ├── geometry.py │ ├── light_transport.py │ ├── raycaster.py │ ├── shader_ops.py │ └── tone_mapping.py ├── version.py └── viewer │ ├── __init__.py │ └── mpl.py ├── docs ├── Makefile ├── make.bat ├── requirements.txt └── source │ ├── .gitignore │ ├── assets │ ├── duffrp-deferred-rast.svg │ ├── game-1024spp-denoised.jpg │ ├── game-8spp-denoised-alpha.jpg │ ├── game-8spp-noisy.jpg │ ├── neural-albedo.png │ ├── procedural-albedo.png │ ├── procedural-normals-small.png │ ├── procedural-normals.png │ ├── spheres-albedo-srgb.png │ ├── spheres-attrs.jpg │ ├── spheres-nvdraa-4xssaa.jpg │ ├── spheres-output.jpg │ ├── spheres-pbr-linear-hdr.jpg │ └── spheres.zip │ ├── concepts.md │ ├── conf.py │ ├── gltf.md │ ├── index.md │ ├── pipeline_comparison.md │ ├── pipeline_surface_deferred.md │ ├── ptpbr.md │ ├── templates │ ├── module.rst_t │ └── package.rst_t │ └── writing_a_material.md └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | tmp 3 | build 4 | dist 5 | *.egg-info 6 | .coverage 7 | .vscode 8 | -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- 1 | version: 2 2 | 3 | build: 4 | os: "ubuntu-22.04" 5 | tools: 6 | python: "3.10" 7 | jobs: 8 | pre_build: 9 | - sphinx-apidoc -eTf -t "docs/source/templates" -o "docs/source/generated" diffrp 10 | 11 | python: 12 | install: 13 | - requirements: docs/requirements.txt 14 | 15 | sphinx: 16 | configuration: docs/source/conf.py 17 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright 2023 eliphat 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Differentiable Render Pipelines `diffrp` 2 | 3 | DiffRP aims to provide an easy-to-use programming interface for **non-differentiable and differentiable** rendering pipelines. 4 | 5 | [[Documentation]](https://diffrp.rtfd.io) 6 | [[Test Suite]](https://github.com/eliphatfs/diffrp-tests) 7 | 8 | ![Teaser](docs/source/assets/game-1024spp-denoised.jpg) 9 | *Rendered with DiffRP.* 10 | 11 | ## Installation 12 | 13 | The package can be installed by: 14 | 15 | ```bash 16 | pip install diffrp 17 | ``` 18 | 19 | ### Note 20 | 21 | DiffRP depends on PyTorch (`torch`). The default version `pip` resolves to may not come with the `cuda` version you want. It is recommended to install [PyTorch](https://pytorch.org/get-started/locally/#start-locally) before you install DiffRP so you can choose the version you like. 22 | 23 | DiffRP rendering is based on CUDA GPUs. You can develop without one, but a CUDA GPU is required to run the code. 24 | 25 | ### Other Dependencies 26 | 27 | If you use rasterization in DiffRP, you need to have the CUDA development kit set up as we use the `nvdiffrast` backend. See also [https://nvlabs.github.io/nvdiffrast/#installation](https://nvlabs.github.io/nvdiffrast/#installation). 28 | 29 | If you want to use hardware ray tracing in DiffRP, you need to install `torchoptix` by `pip install torchoptix`. See also the hardware and driver requirements [https://github.com/eliphatfs/torchoptix?tab=readme-ov-file#requirements](https://github.com/eliphatfs/torchoptix?tab=readme-ov-file#requirements). 30 | 31 | If you plan on using plugins in DiffRP (currently when you compute tangents), `gcc` is required in path. This is already fulfilled in most Linux and Mac distributions. For Windows I recommend the Strawberry Perl [(https://strawberryperl.com/)](https://strawberryperl.com/) distribution of `gcc`. 32 | 33 | ## Get Started 34 | 35 | Rendering attributes of a mesh programmingly is incredibly simple with DiffRP compared to conventional render engines like Blender and Unity. 36 | 37 | In this section, we will get started to render a simple procedural mesh. 38 | 39 | First, let's import what we would need: 40 | 41 | ```python 42 | import diffrp 43 | import trimesh.creation 44 | from diffrp.utils import * 45 | ``` 46 | 47 | We now create a icosphere and render the camera space normal map, where RGB color values in $[0, 1]$ map linearly to $[-1, 1]$ in normal XYZ vectors. 48 | 49 | The first run may take some time due to compiling and importing the `nvdiffrast` backend. Following calls would be fast. 50 | 51 | ```python 52 | # create the mesh (cpu) 53 | mesh = trimesh.creation.icosphere(radius=0.8) 54 | # initialize the DiffRP scene 55 | scene = diffrp.Scene() 56 | # register the mesh, load vertices and faces arrays to GPU 57 | scene.add_mesh_object(diffrp.MeshObject(diffrp.DefaultMaterial(), gpu_f32(mesh.vertices), gpu_i32(mesh.faces))) 58 | # default camera at [0, 0, 3.2] looking backwards 59 | camera = diffrp.PerspectiveCamera() 60 | # create the SurfaceDeferredRenderSession, a deferred-rendering rasterization pipeline session 61 | rp = diffrp.SurfaceDeferredRenderSession(scene, camera) 62 | # convert output tensor to PIL Image and save 63 | to_pil(rp.false_color_camera_space_normal()).save("procedural-normals.png") 64 | ``` 65 | 66 | Only 6 lines of code and we are done! The output should look like this: 67 | 68 | ![Sample](docs/source/assets/procedural-normals-small.png) 69 | 70 | The GPU buffers can be replaced by arbitrary tensors or parameters, and the process has been fully differentiable. 71 | 72 | ## More 73 | 74 | Please refer to the [[Documentation]](https://diffrp.rtfd.io) for more features! 75 | -------------------------------------------------------------------------------- /diffrp/__init__.py: -------------------------------------------------------------------------------- 1 | from .loaders import * 2 | from .materials import * 3 | from .rendering import * 4 | from .resources import * 5 | from .scene import * 6 | from .utils import * 7 | from .utils.cache import * 8 | from .version import __version__ 9 | -------------------------------------------------------------------------------- /diffrp/__main__.py: -------------------------------------------------------------------------------- 1 | from .scripts.main import main 2 | 3 | 4 | if __name__ == '__main__': 5 | main() 6 | -------------------------------------------------------------------------------- /diffrp/loaders/__init__.py: -------------------------------------------------------------------------------- 1 | from .gltf_loader import load_gltf_scene, from_trimesh_scene -------------------------------------------------------------------------------- /diffrp/loaders/gltf_loader.py: -------------------------------------------------------------------------------- 1 | import numpy 2 | import torch 3 | import logging 4 | import trimesh 5 | from PIL import Image 6 | from typing import List, Union, BinaryIO 7 | from trimesh.visual.material import PBRMaterial, SimpleMaterial 8 | from trimesh.visual import ColorVisuals, TextureVisuals 9 | 10 | from ..utils import colors 11 | from ..utils.shader_ops import * 12 | from ..scene import Scene, MeshObject 13 | from ..rendering.camera import PerspectiveCamera 14 | from ..materials.gltf_material import GLTFMaterial, GLTFSampler 15 | from ._trimesh_gltf import load_glb as _load_glb, load_gltf as _load_gltf 16 | 17 | 18 | def ensure_mode(img: Image.Image, mode: str): 19 | if img.mode == mode: 20 | return img 21 | return img.convert(mode) 22 | 23 | 24 | def force_rgba(color: torch.Tensor): 25 | if (color > 1.5).any(): 26 | color = color / 255.0 27 | if color.shape[-1] == 3: 28 | color = float4(color, 1.0) 29 | return color 30 | 31 | 32 | def _simple_to_pbr(mat: SimpleMaterial): 33 | specular = mat.specular.astype(numpy.float32) / 255.0 34 | specular_intensity = specular[0] * 0.2125 + specular[1] * 0.7154 + specular[2] * 0.0721 35 | roughness = (2 / (mat.glossiness + 2)) ** (1.0 / 4.0) 36 | 37 | return PBRMaterial( 38 | roughnessFactor=roughness, 39 | baseColorTexture=mat.image, 40 | baseColorFactor=mat.diffuse, 41 | metallicFactor=specular_intensity, 42 | ) 43 | 44 | 45 | def to_gltf_material(verts: torch.Tensor, visual, material_cache: dict = {}, allow_convert=False): 46 | # GLTF 2.0 specifications 3.9.6 and 5.19 47 | default_mat = GLTFMaterial( 48 | gpu_f32([1, 1, 1, 1]), 49 | GLTFSampler(white_tex()), 50 | 1.0, 1.0, 51 | GLTFSampler(white_tex()), 52 | None, 53 | None, 54 | None, 55 | GLTFSampler(white_tex().rgb), 56 | 0.5, 'OPAQUE' 57 | ) 58 | color = ones_like_vec(verts, 4) 59 | uv = zeros_like_vec(verts, 2) 60 | if isinstance(visual, ColorVisuals): 61 | color = force_rgba(gpu_f32(visual.vertex_colors)) 62 | elif isinstance(visual, TextureVisuals): 63 | if 'color' in visual.vertex_attributes: 64 | color = force_rgba(gpu_f32(visual.vertex_attributes['color'])) 65 | if visual.uv is not None: 66 | uv = gpu_f32(visual.uv) 67 | mat = visual.material 68 | if allow_convert and isinstance(mat, SimpleMaterial): 69 | mat = _simple_to_pbr(mat) 70 | else: 71 | assert isinstance(mat, PBRMaterial), type(mat) 72 | if id(mat) in material_cache: 73 | return uv, color, material_cache[id(mat)] 74 | if mat.baseColorFactor is not None: 75 | default_mat.base_color_factor = force_rgba(gpu_f32(mat.baseColorFactor)) 76 | if mat.baseColorTexture is not None: 77 | base_map = gpu_f32(numpy.asanyarray(ensure_mode(mat.baseColorTexture, "RGBA"))) 78 | base_map /= 255.0 79 | base_map = float4(colors.srgb_to_linear(base_map.rgb), base_map.a) 80 | default_mat.base_color_texture.image = base_map 81 | if mat.metallicFactor is not None: 82 | default_mat.metallic_factor = float(mat.metallicFactor) 83 | if mat.roughnessFactor is not None: 84 | default_mat.roughness_factor = float(mat.roughnessFactor) 85 | if mat.metallicRoughnessTexture is not None: 86 | mr = gpu_f32(numpy.asanyarray(ensure_mode(mat.metallicRoughnessTexture, "RGB"))) 87 | mr /= 255.0 88 | default_mat.metallic_roughness_texture.image = mr 89 | if mat.normalTexture is not None: 90 | nm = gpu_f32(numpy.asanyarray(ensure_mode(mat.normalTexture, "RGB"))) 91 | nm /= 255.0 92 | default_mat.normal_texture = GLTFSampler(nm) 93 | if mat.occlusionTexture is not None: 94 | occ = gpu_f32(numpy.asanyarray(ensure_mode(mat.occlusionTexture, "RGB"))[..., :1]) 95 | occ /= 255.0 96 | default_mat.occlusion_texture = GLTFSampler(occ) 97 | if mat.emissiveFactor is not None and (mat.emissiveFactor > 0).any(): 98 | default_mat.emissive_factor = gpu_f32(mat.emissiveFactor) 99 | if mat.emissiveTexture is not None: 100 | emit = gpu_f32(numpy.asanyarray(ensure_mode(mat.emissiveTexture, "RGB"))) 101 | emit /= 255.0 102 | emit = colors.srgb_to_linear(emit) 103 | default_mat.emissive_texture.image = emit 104 | if mat.alphaCutoff is not None: 105 | default_mat.alpha_cutoff = float(mat.alphaCutoff) 106 | if mat.alphaMode is not None: 107 | default_mat.alpha_mode = str(mat.alphaMode) 108 | material_cache[id(mat)] = default_mat 109 | else: 110 | assert False, ["Unknown visual", type(visual)] 111 | return uv, color, default_mat 112 | 113 | 114 | def from_trimesh_scene(scene: Union[trimesh.Trimesh, trimesh.Scene], compute_tangents=False, allow_convert=True) -> Scene: 115 | """ 116 | Convert a trimesh.Trimesh or trimesh.Scene to a DiffRP Scene. 117 | 118 | Supported metadata: 119 | ``name`` for each mesh object (may be ``None``); 120 | ``camera`` (type PerspectiveCamera) if exists. 121 | 122 | Args: 123 | scene (trimesh.Trimesh | trimesh.Scene): A trimesh loaded scene or mesh. 124 | compute_tangents (bool): 125 | If set, tangents will be computed according to the *MikkTSpace* algorithm. 126 | Execution of the algorithm requires ``gcc`` in the path. 127 | Defaults to ``False``. 128 | Note that computing tangents are not thread-safe. 129 | allow_convert (bool): 130 | DiffRP always try to interpret the materials as PBR materials. 131 | If set, allow a default conversion from simple materials (e.g., from obj and ply formats) 132 | to PBR materials. 133 | 134 | Returns: 135 | The loaded scene. 136 | """ 137 | if isinstance(scene, trimesh.Trimesh): 138 | scene = trimesh.Scene(scene) 139 | drp_scene = Scene() 140 | if scene.has_camera: 141 | camera = PerspectiveCamera( 142 | float(scene.camera.fov[1]), 143 | round(float(scene.camera.resolution[1])), 144 | round(float(scene.camera.resolution[0])), 145 | float(scene.camera.z_near), 146 | float(scene.camera.z_far) 147 | ) 148 | camera.t = scene.camera_transform 149 | drp_scene.metadata['camera'] = camera 150 | meshes: List[trimesh.Trimesh] = [] 151 | transforms: List[torch.Tensor] = [] 152 | names: List[str] = [] 153 | discarded = 0 154 | for node_name in scene.graph.nodes_geometry: 155 | transform, geometry_name = scene.graph[node_name] 156 | # get a copy of the geometry 157 | current = scene.geometry[geometry_name] 158 | if isinstance(current, trimesh.Trimesh): 159 | meshes.append(current) 160 | transforms.append(gpu_f32(transform)) 161 | names.append(geometry_name) 162 | logging.info("Loaded scene with %d submeshes and %d discarded curve/pcd geometry", len(meshes), discarded) 163 | material_cache = {} 164 | for transform, mesh, name in zip(transforms, meshes, names): 165 | verts = gpu_f32(mesh.vertices) 166 | uv, color, mat = to_gltf_material(verts, mesh.visual, material_cache, allow_convert) 167 | # TODO: load vertex tangents if existing 168 | if 'vertex_normals' in mesh._cache and not compute_tangents: 169 | drp_scene.add_mesh_object(MeshObject( 170 | mat, verts, 171 | gpu_i32(mesh.faces), 172 | gpu_f32(mesh.vertex_normals), 173 | transform, color, uv, 174 | metadata=dict(name=name) 175 | )) 176 | else: 177 | flat_idx = mesh.faces.reshape(-1) 178 | # GLTF 2.0 specifications 3.7.2.1 179 | # if normal does not exist in data 180 | # client must calculate *flat* normals 181 | if 'vertex_normals' in mesh._cache: 182 | normals = mesh.vertex_normals[flat_idx] 183 | else: 184 | normals = numpy.stack([mesh.face_normals] * 3, axis=1).reshape(-1, 3) 185 | drp_scene.add_mesh_object(MeshObject( 186 | mat, verts[flat_idx], 187 | gpu_i32(numpy.arange(len(flat_idx)).reshape(-1, 3)), 188 | gpu_f32(normals), 189 | transform, color[flat_idx], uv[flat_idx], 190 | metadata=dict(name=name) 191 | )) 192 | if compute_tangents: 193 | from ..plugins import mikktspace 194 | for mesh_obj in drp_scene.objects: 195 | mesh_obj.tangents = mikktspace.execute(mesh_obj.verts, mesh_obj.normals, mesh_obj.uv) 196 | return drp_scene 197 | 198 | 199 | def load_gltf_scene(path: Union[str, BinaryIO], compute_tangents=False) -> Scene: 200 | """ 201 | Load a glb file as a DiffRP Scene. 202 | 203 | Supported metadata: 204 | ``name`` for each mesh object (may be ``None``); 205 | ``camera`` (type PerspectiveCamera) if exists. 206 | 207 | Args: 208 | path (str | BinaryIO): path to a ``.glb``/``.gltf`` file, or opened ``.glb`` file in binary mode. 209 | compute_tangents (bool): 210 | If set, tangents will be computed according to the *MikkTSpace* algorithm. 211 | Execution of the algorithm requires ``gcc`` in the path. 212 | Defaults to ``False``. 213 | Note that computing tangents are not thread-safe. 214 | 215 | Returns: 216 | The loaded scene. 217 | """ 218 | loader = _load_glb 219 | if isinstance(path, str): 220 | if path.endswith(".gltf"): 221 | loader = _load_gltf 222 | path = open(path, "rb") 223 | kw = loader(path) 224 | path.close() 225 | scene: trimesh.Scene = trimesh.load(kw, force='scene', process=False) 226 | return from_trimesh_scene(scene, compute_tangents, False) 227 | -------------------------------------------------------------------------------- /diffrp/materials/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Includes the interface and built-in implementations of materials. 3 | The APIs in the subpackages are also directly available from this ``diffrp.materials`` package. 4 | """ 5 | 6 | 7 | from .base_material import SurfaceInput, SurfaceUniform, SurfaceOutputStandard, SurfaceMaterial 8 | from .default_material import DefaultMaterial 9 | from .gltf_material import GLTFSampler, GLTFMaterial 10 | -------------------------------------------------------------------------------- /diffrp/materials/base_material.py: -------------------------------------------------------------------------------- 1 | import abc 2 | import torch 3 | from dataclasses import dataclass 4 | from typing import Dict, Optional 5 | from collections.abc import Mapping 6 | from typing_extensions import Literal 7 | from ..utils.cache import cached, key_cached 8 | from ..rendering.interpolator import Interpolator 9 | from ..utils.shader_ops import float4, normalized, transform_point4x3, transform_vector3x3, small_matrix_inverse 10 | 11 | 12 | @dataclass 13 | class VertexArrayObject: 14 | """ 15 | Vertex Array Object. Includes vertex and index buffers of the scene. 16 | 17 | Created automatically by render pipelines. 18 | """ 19 | # geometry 20 | verts: torch.Tensor 21 | normals: torch.Tensor 22 | world_pos: torch.Tensor 23 | 24 | # index buffers 25 | tris: torch.IntTensor 26 | stencils: torch.IntTensor 27 | # REVIEW: this seems coupled with implementation details of render pipelines 28 | # can we make this only for the material? 29 | 30 | # vertex attributes 31 | color: torch.Tensor 32 | uv: torch.Tensor 33 | tangents: torch.Tensor 34 | custom_attrs: Dict[str, torch.Tensor] 35 | 36 | 37 | @dataclass 38 | class SurfaceUniform: 39 | """ 40 | Users can access the system surface uniforms from this class. 41 | 42 | Uniform means the same value or data accessed across a material. 43 | Extra custom uniforms like textures and other attributes should be put in fields of the materials. 44 | 45 | Args: 46 | M (torch.Tensor): The Model matrix. Tensor of shape (4, 4). 47 | V (torch.Tensor): The View matrix in OpenGL convention. Tensor of shape (4, 4). 48 | P (torch.Tensor): The Projection matrix in OpenGL convention. Tensor of shape (4, 4). 49 | """ 50 | M: torch.Tensor 51 | V: torch.Tensor 52 | P: torch.Tensor 53 | 54 | @property 55 | @cached 56 | def camera_matrix(self): 57 | """ 58 | Access the camera matrix, or pose transform, in OpenGL convention. 59 | Camera right, up, forward is X, Y, -Z, respectively. 60 | This is the inverse of the View matrix ``V``. 61 | 62 | Returns: 63 | torch.Tensor: The camera matrix, commonly referred to as ``c2w``. Tensor of shape (4, 4). 64 | """ 65 | return small_matrix_inverse(self.V) 66 | 67 | @property 68 | @cached 69 | def camera_position(self): 70 | """ 71 | Access the position of the camera. 72 | 73 | Returns: 74 | torch.Tensor: The camera position vector XYZ. Tensor of shape (3,). 75 | """ 76 | return self.camera_matrix[:3, 3] 77 | 78 | 79 | class SurfaceInput: 80 | """ 81 | Users can access system and custom attributes of the fragments currently processing 82 | from the ``shade`` function from this class. 83 | """ 84 | 85 | def __init__(self, uniforms: SurfaceUniform, vertex_buffers: VertexArrayObject, interpolator: Interpolator): 86 | self.cache = {} 87 | self.uniforms = uniforms 88 | self.interpolator = interpolator 89 | self.vertex_buffers = vertex_buffers 90 | self.custom_surface_inputs = CustomSurfaceInputs(self) 91 | 92 | def interpolate_ex( 93 | self, 94 | vertex_buffer: torch.Tensor, 95 | world_transform: Literal['none', 'vector', 'vectornor', 'point', 'vector3ex1', 'vector3norex1'] = 'none', 96 | ): 97 | """ 98 | Use this method to interpolate arbitrary vertex attributes into fragments queried in the ``shade`` method. 99 | 100 | Args: 101 | vertex_buffer (torch.Tensor): The vertex attributes to be interpolated. Tensor of shape (V, C). 102 | world_transform (str): 103 | | One of 'none', 'vector', 'vectornor', 'point', 'vector3ex1' or 'vector3norex1'. 104 | | **'none'**: Interpolate the vertex attributes as-is. 105 | | **'vector'**: 106 | Requires ``C = 3``. 107 | Transforms as direction vectors attributed with vertices into world space before interpolating. 108 | | **'vectornor'**: 109 | Requires ``C = 3``. 110 | Transforms as direction vectors attributed with vertices into world space 111 | and normalize before interpolating. 112 | They are **NOT** normalized again after interpolation. 113 | | **'point'**: 114 | Requires ``C = 3``. 115 | Transforms as positions or points attributed with vertices 116 | into world space before interpolating. 117 | | **'vector3ex1'**: 118 | Requires ``C = 4``. 119 | Transforms the first 3 channels as direction vectors attributed with vertices 120 | into world space before interpolating. 121 | The fourth channel is interpolated as-is. 122 | | **'vector3norex1'**: 123 | Requires ``C = 4``. 124 | Transforms the first 3 channels as direction vectors attributed with vertices 125 | into world space and normalize before interpolating. 126 | They are **NOT** normalized again after interpolation. 127 | The fourth channel is interpolated as-is. 128 | 129 | Returns: 130 | torch.Tensor: 131 | Tensor of shape (..., C) where ... is batched pixels. 132 | C is the number of channels in the vertex attribute you input in ``vertex_buffer``. 133 | """ 134 | if world_transform == 'point': 135 | vertex_buffer = transform_point4x3(vertex_buffer, self.uniforms.M) 136 | elif world_transform == 'vector': 137 | vertex_buffer = transform_vector3x3(vertex_buffer, self.uniforms.M) 138 | elif world_transform == 'vectornor': 139 | vertex_buffer = normalized(transform_vector3x3(vertex_buffer, self.uniforms.M)) 140 | elif world_transform == 'vector3ex1': 141 | vertex_buffer = float4(transform_vector3x3(vertex_buffer.xyz, self.uniforms.M), vertex_buffer.w) 142 | elif world_transform == 'vector3norex1': 143 | vertex_buffer = float4(normalized(transform_vector3x3(vertex_buffer.xyz, self.uniforms.M)), vertex_buffer.w) 144 | return self.interpolator.interpolate(vertex_buffer) 145 | 146 | @property 147 | @cached 148 | def view_dir(self) -> torch.Tensor: 149 | """ 150 | View directions. 151 | The normalized unit vectors looking from the camera to the world position of the fragments. 152 | 153 | Returns: 154 | torch.Tensor: 155 | Tensor of shape (..., 3) where ... is batched pixels. 156 | """ 157 | return normalized(self.world_pos - self.uniforms.camera_position) 158 | 159 | @property 160 | @cached 161 | def world_pos(self) -> torch.Tensor: 162 | """ 163 | The world position of the fragments. 164 | 165 | Returns: 166 | torch.Tensor: 167 | Tensor of shape (..., 3) where ... is batched pixels. 168 | """ 169 | return self.interpolate_ex(self.vertex_buffers.world_pos) 170 | 171 | @property 172 | @cached 173 | def local_pos(self) -> torch.Tensor: 174 | """ 175 | The local (object space) position of the fragments. 176 | 177 | Returns: 178 | torch.Tensor: 179 | Tensor of shape (..., 3) where ... is batched pixels. 180 | """ 181 | return self.interpolate_ex(self.vertex_buffers.verts) 182 | 183 | @property 184 | @cached 185 | def world_normal_unnormalized(self) -> torch.Tensor: 186 | """ 187 | The geometry world normal of the fragments. 188 | This includes the effect of smooth (vertex interpolation) and flat face specification in the ``MeshObject``. 189 | 190 | Normals are normalized in vertices, but not again after interpolation. 191 | Thus, this is slightly off from unit vectors, but may be useful in e.g. tangent space calculations. 192 | 193 | Returns: 194 | torch.Tensor: 195 | Tensor of shape (..., 3) where ... is batched pixels. 196 | """ 197 | return self.interpolate_ex(self.vertex_buffers.normals, 'vectornor') 198 | 199 | @property 200 | @cached 201 | def world_normal(self) -> torch.Tensor: 202 | """ 203 | The normalized unit vectors for geometry world normal of the fragments. 204 | This includes the effect of smooth (vertex interpolation) and flat face specification in the ``MeshObject``. 205 | 206 | Returns: 207 | torch.Tensor: 208 | Tensor of shape (..., 3) where ... is batched pixels. 209 | """ 210 | return normalized(self.world_normal_unnormalized) 211 | 212 | @property 213 | @cached 214 | def color(self) -> torch.Tensor: 215 | """ 216 | The interpolated color attributes (RGBA) of the fragments. 217 | Usually ranges from 0 to 1 (instead of 255). 218 | 219 | Defaults to white (ones) if not specified in the ``MeshObject``. 220 | 221 | Returns: 222 | torch.Tensor: 223 | Tensor of shape (..., 4) where ... is batched pixels. 224 | """ 225 | # F4, vertex color, default to ones 226 | return self.interpolate_ex(self.vertex_buffers.color) 227 | 228 | @property 229 | @cached 230 | def uv(self) -> torch.Tensor: 231 | """ 232 | The interpolated UV coordinate attributes of the fragments. 233 | (0, 0) is the bottom-left in OpenGL convention, in contrast to top-left in most image libraries. 234 | 235 | Defaults to zeros if not specified in the ``MeshObject``. 236 | 237 | Returns: 238 | torch.Tensor: 239 | Tensor of shape (..., 2) where ... is batched pixels. 240 | """ 241 | return self.interpolate_ex(self.vertex_buffers.uv) 242 | 243 | @property 244 | @cached 245 | def world_tangent(self) -> torch.Tensor: 246 | """ 247 | The interpolated tangents in world space of the fragments. 248 | 249 | The xyz dimensions form the tangent vector, while the w dimension specifies the sign (1 or -1). 250 | 251 | Defaults to zeros if not specified in the ``MeshObject``. 252 | 253 | Returns: 254 | torch.Tensor: 255 | Tensor of shape (..., 4) where ... is batched pixels. 256 | """ 257 | return self.interpolate_ex(self.vertex_buffers.tangents, 'vector3norex1') 258 | 259 | @property 260 | def custom_attrs(self) -> Mapping[str, torch.Tensor]: 261 | """ 262 | Dictionary of custom attributes of the fragments. 263 | 264 | The keys are the same as you specified in the ``MeshObject``. 265 | 266 | If some objects provide the custom inputs while others not for the key, 267 | you will get zeros for objects where the custom vertex attributes are absent. 268 | 269 | An error will be raised if custom inputs with the same key have different numbers of channels. 270 | 271 | Returns: 272 | torch.Tensor: 273 | Tensor of shape (..., C) where ... is batched pixels. 274 | C is the number of channels in the vertex attribute you input in ``MeshObject``. 275 | """ 276 | return self.custom_surface_inputs 277 | 278 | 279 | class CustomSurfaceInputs(Mapping): 280 | """ 281 | Implementation for custom vertex attribute access. 282 | 283 | This is an implementation detail and you do not need to take care of it. 284 | """ 285 | 286 | def __init__(self, si: SurfaceInput) -> None: 287 | self.si = si 288 | 289 | @key_cached 290 | def __getitem__(self, attr: str) -> torch.Tensor: 291 | return self.si.interpolate_ex(self.si.vertex_buffers.custom_attrs[attr]) 292 | 293 | def __len__(self) -> int: 294 | return len(self.si.vertex_buffers.custom_attrs) 295 | 296 | def __iter__(self) -> iter: 297 | return iter(self.si.vertex_buffers.custom_attrs) 298 | 299 | 300 | @dataclass 301 | class SurfaceOutputStandard: 302 | """ 303 | The output structure that specifies standard and custom (AOV) outputs from the material (shader). 304 | All outputs are optional. 305 | The outputs shapes are all (..., C), where (...) is batched pixels. 306 | They should match the batch dimensions provided by ``SurfaceInput``. 307 | It can be 1D to 3D according to interpolator implementation, 1D under default settings 308 | -- a batch of pixel features of shape (B, C). 309 | See the documentation for interpolator options and interpolators for more details. 310 | 311 | Args: 312 | albedo (torch.Tensor): 313 | The per-pixel base color in linear RGB space, used both for diffuse and specular. 314 | Defaults to magenta (1.0, 0.0, 1.0). 315 | Tensor of shape (..., 3), where (...) should match the pixel batch dimensions in inputs. 316 | normal (torch.Tensor): 317 | The per-pixel normal in specified space. Defaults to geometry normal. 318 | See also ``normal_space``. 319 | This can be used to implement per-pixel (neural) normal maps. 320 | Tensor of shape (..., 3), where (...) should match the pixel batch dimensions in inputs. 321 | emission (torch.Tensor): 322 | The per-pixel emission value in linear RGB space. 323 | Defaults to black (0.0, 0.0, 0.0) which means no emission. 324 | Tensor of shape (..., 3), where (...) should match the pixel batch dimensions in inputs. 325 | metallic (torch.Tensor): 326 | The per-pixel metallic value. 327 | Defaults to 0.0, fully dielectric. 328 | Tensor of shape (..., 1), where (...) should match the pixel batch dimensions in inputs. 329 | smoothness (torch.Tensor): 330 | The per-pixel smoothness value. 331 | Defaults to 0.5. 332 | Tensor of shape (..., 1), where (...) should match the pixel batch dimensions in inputs. 333 | occlusion (torch.Tensor): 334 | The per-pixel ambient occlusion value. Dims indirect light. 335 | Defaults to 1.0, no occlusion. 336 | Tensor of shape (..., 1), where (...) should match the pixel batch dimensions in inputs. 337 | alpha (torch.Tensor): 338 | The per-pixel alpha value. Linear transparency. 339 | Defaults to 1.0, fully opaque. 340 | Any value will be overridden to 1.0 if render sessions are executed in ``opaque_only`` mode. 341 | Tensor of shape (..., 1), where (...) should match the pixel batch dimensions in inputs. 342 | aovs (Dict[str, torch.Tensor]): 343 | Auxiliary Output Variables (AOVs). 344 | Anything you would like to output for the shader. 345 | Batch dimensions of values need to match the pixel batch dimensions in inputs. 346 | normal_space(str): 347 | | One of 'tangent', 'object' and 'world'. 348 | | Defaults to 'tangent' as most normal maps are in tangent space, 349 | featuring a 'blueish' look. 350 | | 'object' and 'world' spaces are simple. 351 | | 'tangent' means (tangent, bitangent, geometry_normal) for XYZ components, respectively, 352 | where tangents are typically generated from the UV mapping. 353 | See also the ``mikktspace`` plugin. 354 | """ 355 | albedo: Optional[torch.Tensor] = None # F3, base color (diffuse or specular), default to magenta 356 | normal: Optional[torch.Tensor] = None # F3, normal in specified space, default to geometry normal 357 | emission: Optional[torch.Tensor] = None # F3, default to black 358 | metallic: Optional[torch.Tensor] = None # F1, default to 0.0 359 | smoothness: Optional[torch.Tensor] = None # F1, default to 0.5 360 | occlusion: Optional[torch.Tensor] = None # F1, default to 1.0 361 | alpha: Optional[torch.Tensor] = None # F1, default to 1.0 362 | aovs: Optional[Dict[str, torch.Tensor]] = None # arbitrary, default to 0.0 for non-existing areas in render 363 | normal_space: Literal['tangent', 'object', 'world'] = 'tangent' 364 | 365 | 366 | class SurfaceMaterial(metaclass=abc.ABCMeta): 367 | 368 | @abc.abstractmethod 369 | def shade(self, su: SurfaceUniform, si: SurfaceInput) -> SurfaceOutputStandard: 370 | """ 371 | The interface for implementing materials. 372 | It takes a fragment batch specified as ``SurfaceUniform`` and ``SurfaceInput``. 373 | 374 | You shall not assume any order or shape of the fragment batch. 375 | In other words, the method should be trivial on the batch dimension, that is, 376 | the method should return batched outputs equivalent to 377 | concatenated outputs from the same batch of input but split into multiple sub-batches. 378 | 379 | Most functions implemented with attribute accesses, element/vector-wise operations 380 | and DiffRP utility functions naturally have this property. 381 | """ 382 | raise NotImplementedError 383 | -------------------------------------------------------------------------------- /diffrp/materials/default_material.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Optional 3 | from .base_material import SurfaceInput, SurfaceMaterial, SurfaceUniform, SurfaceOutputStandard 4 | 5 | 6 | class DefaultMaterial(SurfaceMaterial): 7 | """ 8 | Default material. 9 | 10 | Outputs the interpolated vertex color attributes combined with an optional tint as the albedo, 11 | and other attributes remain the default. 12 | 13 | Args: 14 | tint(Optional[torch.Tensor]): Linear multipliers for colors. Defaults to ``None``. 15 | """ 16 | 17 | def __init__(self, tint: Optional[torch.Tensor] = None) -> None: 18 | super().__init__() 19 | self.tint = tint 20 | 21 | def shade(self, su: SurfaceUniform, si: SurfaceInput) -> SurfaceOutputStandard: 22 | if self.tint is not None: 23 | return SurfaceOutputStandard(si.color.rgb * self.tint) 24 | return SurfaceOutputStandard(si.color.rgb) 25 | -------------------------------------------------------------------------------- /diffrp/materials/gltf_material.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Optional 3 | from dataclasses import dataclass 4 | from typing_extensions import Literal 5 | from ..utils.shader_ops import sample2d, fsa, ones_like_vec 6 | from .base_material import SurfaceInput, SurfaceMaterial, SurfaceOutputStandard, SurfaceUniform 7 | 8 | 9 | @dataclass 10 | class GLTFSampler: 11 | image: torch.Tensor # F1, F3 or F4 12 | wrap_mode: Literal['repeat', 'clamp', 'mirror'] = 'repeat' 13 | interpolation: Literal['point', 'linear'] = 'linear' 14 | 15 | def sample(self, uv: torch.Tensor): 16 | if self.wrap_mode == 'repeat': 17 | uv = uv.remainder(1.0) 18 | return sample2d( 19 | self.image, uv, 20 | wrap='border' if self.wrap_mode == 'clamp' else 'reflection', 21 | mode='bilinear' if self.interpolation == 'linear' else 'nearest' 22 | ) 23 | 24 | 25 | @dataclass 26 | class GLTFMaterial(SurfaceMaterial): 27 | """ 28 | A standard PBR material coherent with GLTF 2.0 Specifications. 29 | """ 30 | base_color_factor: torch.Tensor # F4 31 | base_color_texture: GLTFSampler 32 | 33 | metallic_factor: float 34 | roughness_factor: float 35 | metallic_roughness_texture: GLTFSampler 36 | 37 | normal_texture: Optional[GLTFSampler] 38 | occlusion_texture: Optional[GLTFSampler] 39 | 40 | emissive_factor: Optional[torch.Tensor] # F3 41 | emissive_texture: GLTFSampler 42 | 43 | alpha_cutoff: float 44 | alpha_mode: Literal['OPAQUE', 'MASK', 'BLEND'] 45 | 46 | def shade(self, su: SurfaceUniform, si: SurfaceInput) -> SurfaceOutputStandard: 47 | rgba = self.base_color_factor * si.color * self.base_color_texture.sample(si.uv) 48 | mr = self.metallic_roughness_texture.sample(si.uv) 49 | 50 | if self.alpha_mode == 'OPAQUE': 51 | alpha = None 52 | elif self.alpha_mode == 'MASK': 53 | alpha = (rgba.a > self.alpha_cutoff).float() 54 | elif self.alpha_mode == 'BLEND': 55 | alpha = rgba.a 56 | else: 57 | raise ValueError('bad GLTFMaterial.alpha_mode', self.alpha_mode) 58 | 59 | return SurfaceOutputStandard( 60 | rgba.rgb, 61 | fsa(2, self.normal_texture.sample(si.uv), -1) if self.normal_texture is not None else None, 62 | self.emissive_factor * self.emissive_texture.sample(si.uv) if self.emissive_factor is not None else None, 63 | self.metallic_factor * mr.b, 64 | fsa(-self.roughness_factor, mr.g, 1.0), 65 | self.occlusion_texture.sample(si.uv).r if self.occlusion_texture is not None else ones_like_vec(si.uv, 1), 66 | alpha 67 | ) 68 | -------------------------------------------------------------------------------- /diffrp/plugins/.gitignore: -------------------------------------------------------------------------------- 1 | *.bin 2 | -------------------------------------------------------------------------------- /diffrp/plugins/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/diffrp/plugins/__init__.py -------------------------------------------------------------------------------- /diffrp/plugins/mikktspace/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import ctypes 4 | import subprocess 5 | 6 | 7 | wd = os.path.dirname(os.path.abspath(__file__)) 8 | plugin_name = 'mikktspace_v1.bin' 9 | 10 | 11 | def compile_plugin(): 12 | subprocess.check_call([ 13 | "gcc", 14 | "mikktinterface.c", "mikktspace.c", 15 | "-shared", "-fPIC", "-Ofast", 16 | "-o", plugin_name 17 | ], cwd=wd) 18 | 19 | 20 | if not os.path.exists(os.path.join(wd, plugin_name)): 21 | compile_plugin() 22 | 23 | 24 | mikkt = ctypes.CDLL(os.path.join(wd, plugin_name)) 25 | mikkt.compute_tri_tangents.argtypes = (ctypes.c_void_p, ctypes.c_void_p, ctypes.c_longlong) 26 | mikkt.compute_tri_tangents.restype = None 27 | 28 | 29 | @torch.no_grad() 30 | def execute(verts: torch.Tensor, normals: torch.Tensor, uv: torch.Tensor): 31 | """ 32 | Execute the *MikkTSpace* algorithm to 33 | compute tangent vectors and bitangent signs, which form the tangent space. 34 | The algorithm is serial algorithm on CPU and not differentiable, 35 | but tensors are automatically transferred to and back from CPU. 36 | 37 | It takes flat, unrolled vertex positions, normals and uvs. 38 | 39 | The plugin is not thread-safe. 40 | 41 | Args: 42 | verts (torch.Tensor): (F * 3, 3) vertex positions of faces unrolled. 43 | normals (torch.Tensor): (F * 3, 3) vertex normals of faces unrolled. 44 | uv (torch.Tensor): (F * 3, 2) vertex UV coordinates of faces unrolled. 45 | 46 | Returns: 47 | torch.Tensor: (F * 3, 4) tangent space. Has the same device and dtype as ``verts``. 48 | """ 49 | vbuf = torch.cat([verts, normals, uv], dim=-1).to('cpu', torch.float32).contiguous() 50 | tangent = torch.empty([len(verts), 4], dtype=torch.float32, memory_format=torch.contiguous_format) 51 | mikkt.compute_tri_tangents(vbuf.data_ptr(), tangent.data_ptr(), len(vbuf)) 52 | return tangent.to(verts) 53 | -------------------------------------------------------------------------------- /diffrp/plugins/mikktspace/mikktinterface.c: -------------------------------------------------------------------------------- 1 | // 2 | // Modified from mikktpy to support buffer interface 3 | // Original Copyright (c) 2013 Ambrus Csaszar 4 | // 5 | #include 6 | #include "mikktspace.h" 7 | 8 | 9 | static const int components_per_vert = 3 + 3 + 2; // xyz, normal-xyz, uv 10 | static const int components_per_tangent = 4; // tangent + sign 11 | static const int vertices_per_face = 3; 12 | 13 | typedef struct 14 | { 15 | const float * v; 16 | float * r; 17 | long long n; 18 | } TriTangentState; 19 | 20 | // Returns the number of faces (triangles/quads) on the mesh to be processed. 21 | int getNumFaces(const SMikkTSpaceContext * pContext); 22 | 23 | // Returns the number of vertices on face number iFace 24 | // iFace is a number in the range {0, 1, ..., getNumFaces()-1} 25 | int getNumVerticesOfFace(const SMikkTSpaceContext * pContext, const int iFace); 26 | 27 | // returns the position/normal/texcoord of the referenced face of vertex number iVert. 28 | // iVert is in the range {0,1,2} for triangles and {0,1,2,3} for quads. 29 | void getPosition(const SMikkTSpaceContext * pContext, float fvPosOut[], const int iFace, const int iVert); 30 | void getNormal(const SMikkTSpaceContext * pContext, float fvNormOut[], const int iFace, const int iVert); 31 | void getTexCoord(const SMikkTSpaceContext * pContext, float fvTexcOut[], const int iFace, const int iVert); 32 | 33 | // either (or both) of the two setTSpace callbacks can be set. 34 | // The call-back m_setTSpaceBasic() is sufficient for basic normal mapping. 35 | 36 | // This function is used to return the tangent and fSign to the application. 37 | // fvTangent is a unit length vector. 38 | // For normal maps it is sufficient to use the following simplified version of the bitangent which is generated at pixel/vertex level. 39 | // bitangent = fSign * cross(vN, tangent); 40 | // Note that the results are returned unindexed. It is possible to generate a new index list 41 | // But averaging/overwriting tangent spaces by using an already existing index list WILL produce INCRORRECT results. 42 | // DO NOT! use an already existing index list. 43 | void setTSpaceBasic(const SMikkTSpaceContext * pContext, const float fvTangent[], const float fSign, const int iFace, const int iVert); 44 | 45 | 46 | int getNumFaces(const SMikkTSpaceContext * pContext) 47 | { 48 | TriTangentState *s = (TriTangentState*)pContext->m_pUserData; 49 | int ret = s->n / vertices_per_face; 50 | // cout << "GetNumFaces: " << ret << endl; 51 | return ret; 52 | } 53 | 54 | int getNumVerticesOfFace(const SMikkTSpaceContext * pContext, const int iFace) 55 | { 56 | // cout << "GetVerticesPerFace: " << vertices_per_face << endl; 57 | return vertices_per_face; 58 | } 59 | 60 | void getPosition(const SMikkTSpaceContext * pContext, float fvPosOut[], const int iFace, const int iVert) 61 | { 62 | TriTangentState *s = (TriTangentState*)pContext->m_pUserData; 63 | int idx = iFace * vertices_per_face * components_per_vert + iVert * components_per_vert; 64 | // cout << "GetPosition: " << idx << endl; 65 | fvPosOut[0] = s->v[idx + 0]; 66 | fvPosOut[1] = s->v[idx + 1]; 67 | fvPosOut[2] = s->v[idx + 2]; 68 | } 69 | 70 | void getNormal(const SMikkTSpaceContext * pContext, float fvNormOut[], const int iFace, const int iVert) 71 | { 72 | TriTangentState *s = (TriTangentState*)pContext->m_pUserData; 73 | int idx = iFace * vertices_per_face * components_per_vert + iVert * components_per_vert; 74 | fvNormOut[0] = s->v[idx + 3]; 75 | fvNormOut[1] = s->v[idx + 4]; 76 | fvNormOut[2] = s->v[idx + 5]; 77 | } 78 | 79 | void getTexCoord(const SMikkTSpaceContext * pContext, float fvTexcOut[], const int iFace, const int iVert) 80 | { 81 | TriTangentState *s = (TriTangentState*)pContext->m_pUserData; 82 | int idx = iFace * vertices_per_face * components_per_vert + iVert * components_per_vert; 83 | fvTexcOut[0] = s->v[idx + 6]; 84 | fvTexcOut[1] = s->v[idx + 7]; 85 | } 86 | 87 | void setTSpaceBasic(const SMikkTSpaceContext * pContext, const float fvTangent[], const float fSign, const int iFace, const int iVert) 88 | { 89 | TriTangentState *s = (TriTangentState*)pContext->m_pUserData; 90 | int idx = iFace * vertices_per_face * components_per_tangent + iVert * components_per_tangent; 91 | 92 | // cout << "SetTspace (" << iFace << "," << iVert << ") "; 93 | // for (int i = 0; i < 3; i++) 94 | // cout << (i>0 ? "," : "") << fvTangent[i]; 95 | // cout << "; "; 96 | // for (int i = 0; i < 3; i++) 97 | // cout << (i>0 ? "," : "") << fvBiTangent[i]; 98 | // cout << endl; 99 | // cout << "idx: " << idx << endl; 100 | 101 | s->r[idx + 0] = fvTangent[0]; 102 | s->r[idx + 1] = fvTangent[1]; 103 | s->r[idx + 2] = fvTangent[2]; 104 | s->r[idx + 3] = fSign; 105 | } 106 | 107 | void compute_tri_tangents(const float * v, float * r, long long n) 108 | { 109 | SMikkTSpaceInterface mikktInterface; 110 | mikktInterface.m_getNumFaces = getNumFaces; 111 | mikktInterface.m_getNumVerticesOfFace = getNumVerticesOfFace; 112 | mikktInterface.m_getPosition = getPosition; 113 | mikktInterface.m_getNormal = getNormal; 114 | mikktInterface.m_getTexCoord = getTexCoord; 115 | mikktInterface.m_setTSpaceBasic = setTSpaceBasic; 116 | mikktInterface.m_setTSpace = NULL; 117 | 118 | TriTangentState s = { v, r, n }; 119 | 120 | SMikkTSpaceContext mikktContext; 121 | mikktContext.m_pInterface = &mikktInterface; 122 | mikktContext.m_pUserData = &s; 123 | 124 | genTangSpaceDefault(&mikktContext); 125 | } 126 | -------------------------------------------------------------------------------- /diffrp/plugins/mikktspace/mikktspace.h: -------------------------------------------------------------------------------- 1 | /** \file mikktspace/mikktspace.h 2 | * \ingroup mikktspace 3 | */ 4 | /** 5 | * Copyright (C) 2011 by Morten S. Mikkelsen 6 | * 7 | * This software is provided 'as-is', without any express or implied 8 | * warranty. In no event will the authors be held liable for any damages 9 | * arising from the use of this software. 10 | * 11 | * Permission is granted to anyone to use this software for any purpose, 12 | * including commercial applications, and to alter it and redistribute it 13 | * freely, subject to the following restrictions: 14 | * 15 | * 1. The origin of this software must not be misrepresented; you must not 16 | * claim that you wrote the original software. If you use this software 17 | * in a product, an acknowledgment in the product documentation would be 18 | * appreciated but is not required. 19 | * 2. Altered source versions must be plainly marked as such, and must not be 20 | * misrepresented as being the original software. 21 | * 3. This notice may not be removed or altered from any source distribution. 22 | */ 23 | 24 | #ifndef __MIKKTSPACE_H__ 25 | #define __MIKKTSPACE_H__ 26 | 27 | 28 | #ifdef __cplusplus 29 | extern "C" { 30 | #endif 31 | 32 | /* Author: Morten S. Mikkelsen 33 | * Version: 1.0 34 | * 35 | * The files mikktspace.h and mikktspace.c are designed to be 36 | * stand-alone files and it is important that they are kept this way. 37 | * Not having dependencies on structures/classes/libraries specific 38 | * to the program, in which they are used, allows them to be copied 39 | * and used as is into any tool, program or plugin. 40 | * The code is designed to consistently generate the same 41 | * tangent spaces, for a given mesh, in any tool in which it is used. 42 | * This is done by performing an internal welding step and subsequently an order-independent evaluation 43 | * of tangent space for meshes consisting of triangles and quads. 44 | * This means faces can be received in any order and the same is true for 45 | * the order of vertices of each face. The generated result will not be affected 46 | * by such reordering. Additionally, whether degenerate (vertices or texture coordinates) 47 | * primitives are present or not will not affect the generated results either. 48 | * Once tangent space calculation is done the vertices of degenerate primitives will simply 49 | * inherit tangent space from neighboring non degenerate primitives. 50 | * The analysis behind this implementation can be found in my master's thesis 51 | * which is available for download --> http://image.diku.dk/projects/media/morten.mikkelsen.08.pdf 52 | * Note that though the tangent spaces at the vertices are generated in an order-independent way, 53 | * by this implementation, the interpolated tangent space is still affected by which diagonal is 54 | * chosen to split each quad. A sensible solution is to have your tools pipeline always 55 | * split quads by the shortest diagonal. This choice is order-independent and works with mirroring. 56 | * If these have the same length then compare the diagonals defined by the texture coordinates. 57 | * XNormal which is a tool for baking normal maps allows you to write your own tangent space plugin 58 | * and also quad triangulator plugin. 59 | */ 60 | 61 | 62 | typedef int tbool; 63 | typedef struct SMikkTSpaceContext SMikkTSpaceContext; 64 | 65 | typedef struct { 66 | // Returns the number of faces (triangles/quads) on the mesh to be processed. 67 | int (*m_getNumFaces)(const SMikkTSpaceContext * pContext); 68 | 69 | // Returns the number of vertices on face number iFace 70 | // iFace is a number in the range {0, 1, ..., getNumFaces()-1} 71 | int (*m_getNumVerticesOfFace)(const SMikkTSpaceContext * pContext, const int iFace); 72 | 73 | // returns the position/normal/texcoord of the referenced face of vertex number iVert. 74 | // iVert is in the range {0,1,2} for triangles and {0,1,2,3} for quads. 75 | void (*m_getPosition)(const SMikkTSpaceContext * pContext, float fvPosOut[], const int iFace, const int iVert); 76 | void (*m_getNormal)(const SMikkTSpaceContext * pContext, float fvNormOut[], const int iFace, const int iVert); 77 | void (*m_getTexCoord)(const SMikkTSpaceContext * pContext, float fvTexcOut[], const int iFace, const int iVert); 78 | 79 | // either (or both) of the two setTSpace callbacks can be set. 80 | // The call-back m_setTSpaceBasic() is sufficient for basic normal mapping. 81 | 82 | // This function is used to return the tangent and fSign to the application. 83 | // fvTangent is a unit length vector. 84 | // For normal maps it is sufficient to use the following simplified version of the bitangent which is generated at pixel/vertex level. 85 | // bitangent = fSign * cross(vN, tangent); 86 | // Note that the results are returned unindexed. It is possible to generate a new index list 87 | // But averaging/overwriting tangent spaces by using an already existing index list WILL produce INCRORRECT results. 88 | // DO NOT! use an already existing index list. 89 | void (*m_setTSpaceBasic)(const SMikkTSpaceContext * pContext, const float fvTangent[], const float fSign, const int iFace, const int iVert); 90 | 91 | // This function is used to return tangent space results to the application. 92 | // fvTangent and fvBiTangent are unit length vectors and fMagS and fMagT are their 93 | // true magnitudes which can be used for relief mapping effects. 94 | // fvBiTangent is the "real" bitangent and thus may not be perpendicular to fvTangent. 95 | // However, both are perpendicular to the vertex normal. 96 | // For normal maps it is sufficient to use the following simplified version of the bitangent which is generated at pixel/vertex level. 97 | // fSign = bIsOrientationPreserving ? 1.0f : (-1.0f); 98 | // bitangent = fSign * cross(vN, tangent); 99 | // Note that the results are returned unindexed. It is possible to generate a new index list 100 | // But averaging/overwriting tangent spaces by using an already existing index list WILL produce INCRORRECT results. 101 | // DO NOT! use an already existing index list. 102 | void (*m_setTSpace)(const SMikkTSpaceContext * pContext, const float fvTangent[], const float fvBiTangent[], const float fMagS, const float fMagT, 103 | const tbool bIsOrientationPreserving, const int iFace, const int iVert); 104 | } SMikkTSpaceInterface; 105 | 106 | struct SMikkTSpaceContext 107 | { 108 | SMikkTSpaceInterface * m_pInterface; // initialized with callback functions 109 | void * m_pUserData; // pointer to client side mesh data etc. (passed as the first parameter with every interface call) 110 | }; 111 | 112 | // these are both thread safe! 113 | tbool genTangSpaceDefault(const SMikkTSpaceContext * pContext); // Default (recommended) fAngularThreshold is 180 degrees (which means threshold disabled) 114 | tbool genTangSpace(const SMikkTSpaceContext * pContext, const float fAngularThreshold); 115 | 116 | 117 | // To avoid visual errors (distortions/unwanted hard edges in lighting), when using sampled normal maps, the 118 | // normal map sampler must use the exact inverse of the pixel shader transformation. 119 | // The most efficient transformation we can possibly do in the pixel shader is 120 | // achieved by using, directly, the "unnormalized" interpolated tangent, bitangent and vertex normal: vT, vB and vN. 121 | // pixel shader (fast transform out) 122 | // vNout = normalize( vNt.x * vT + vNt.y * vB + vNt.z * vN ); 123 | // where vNt is the tangent space normal. The normal map sampler must likewise use the 124 | // interpolated and "unnormalized" tangent, bitangent and vertex normal to be compliant with the pixel shader. 125 | // sampler does (exact inverse of pixel shader): 126 | // float3 row0 = cross(vB, vN); 127 | // float3 row1 = cross(vN, vT); 128 | // float3 row2 = cross(vT, vB); 129 | // float fSign = dot(vT, row0)<0 ? -1 : 1; 130 | // vNt = normalize( fSign * float3(dot(vNout,row0), dot(vNout,row1), dot(vNout,row2)) ); 131 | // where vNout is the sampled normal in some chosen 3D space. 132 | // 133 | // Should you choose to reconstruct the bitangent in the pixel shader instead 134 | // of the vertex shader, as explained earlier, then be sure to do this in the normal map sampler also. 135 | // Finally, beware of quad triangulations. If the normal map sampler doesn't use the same triangulation of 136 | // quads as your renderer then problems will occur since the interpolated tangent spaces will differ 137 | // eventhough the vertex level tangent spaces match. This can be solved either by triangulating before 138 | // sampling/exporting or by using the order-independent choice of diagonal for splitting quads suggested earlier. 139 | // However, this must be used both by the sampler and your tools/rendering pipeline. 140 | 141 | #ifdef __cplusplus 142 | } 143 | #endif 144 | 145 | #endif 146 | -------------------------------------------------------------------------------- /diffrp/rendering/__init__.py: -------------------------------------------------------------------------------- 1 | from .camera import Camera, RawCamera, PerspectiveCamera 2 | from .denoiser import get_denoiser, run_denoiser 3 | from .interpolator import Interpolator, MaskedSparseInterpolator, FullScreenInterpolator, polyfill_interpolate 4 | from .mixin import RenderSessionMixin 5 | from .path_tracing import PathTracingSession, PathTracingSessionOptions, RayOutputs 6 | from .surface_deferred import SurfaceDeferredRenderSession, SurfaceDeferredRenderSessionOptions 7 | -------------------------------------------------------------------------------- /diffrp/rendering/camera.py: -------------------------------------------------------------------------------- 1 | import math 2 | import numpy 3 | import torch 4 | import trimesh 5 | import calibur 6 | from typing import Union, List 7 | from ..utils.shader_ops import gpu_f32 8 | 9 | 10 | class Camera: 11 | """ 12 | Abstract class for camera specification. 13 | 14 | You need to implement the ``V()``, ``P()`` and ``resolution()`` methods 15 | if you inherit from the class. 16 | """ 17 | 18 | def __init__(self) -> None: 19 | self.t = trimesh.transformations.translation_matrix([0, 0.0, 3.2]) 20 | 21 | def V(self): 22 | """ 23 | Returns: 24 | torch.Tensor: 25 | | GPU tensor of shape (4, 4). GL view matrix. 26 | | Inverse of the c2w or camera pose transform matrix. 27 | | You may need to convert the camera poses if the pose is not in GL convention 28 | (X right, Y up, -Z forward). You may use ``calibur`` for this. 29 | | It should be an affine transform matrix with the last row fixed to [0, 0, 0, 1]. 30 | It is also often called an extrinsic matrix or w2c matrix in other conventions. 31 | """ 32 | return gpu_f32(trimesh.transformations.inverse_matrix(self.t).astype(numpy.float32)) 33 | 34 | def P(self): 35 | """ 36 | Returns: 37 | torch.Tensor: GPU tensor of shape (4, 4). GL projection matrix. 38 | """ 39 | raise NotImplementedError 40 | 41 | def resolution(self): 42 | """ 43 | Returns: 44 | Tuple[int, int]: Resolution in (h, w) order. 45 | """ 46 | raise NotImplementedError 47 | 48 | 49 | class RawCamera(Camera): 50 | """ 51 | Raw data-driven camera that takes GPU tensors of the view and projection matrices to drive the camera. 52 | """ 53 | def __init__(self, h: int, w: int, v: torch.Tensor, p: torch.Tensor) -> None: 54 | self.h = h 55 | self.w = w 56 | self.v = v 57 | self.p = p 58 | 59 | def V(self): 60 | return self.v 61 | 62 | def P(self): 63 | return self.p 64 | 65 | def resolution(self): 66 | return self.h, self.w 67 | 68 | class PerspectiveCamera(Camera): 69 | """ 70 | Perspective camera class. Angles are in degrees. 71 | """ 72 | def __init__(self, fov=30, h=512, w=512, near=0.1, far=10.0) -> None: 73 | super().__init__() 74 | self.fov = fov 75 | self.h = h 76 | self.w = w 77 | self.near = near 78 | self.far = far 79 | 80 | def P(self): 81 | cx, cy = self.w / 2, self.h / 2 82 | fx = fy = calibur.fov_to_focal(numpy.radians(self.fov), self.h) 83 | return gpu_f32(calibur.projection_gl_persp( 84 | self.w, self.h, 85 | cx, cy, 86 | fx, fy, 87 | self.near, self.far 88 | )) 89 | 90 | def resolution(self): 91 | return self.h, self.w 92 | 93 | def set_transform(self, tr: numpy.ndarray): 94 | self.t = tr 95 | 96 | def lookat(self, point: Union[List[int], numpy.ndarray]): 97 | if isinstance(point, list): 98 | point = numpy.array(point) 99 | assert list(point.shape) == [3] 100 | pos = trimesh.transformations.translation_from_matrix(self.t) 101 | fwd = point - pos 102 | z = -fwd 103 | x = numpy.cross(fwd, [0, 1, 0]) 104 | y = numpy.cross(x, fwd) 105 | self.t[:3, 0] = calibur.normalized(x) 106 | self.t[:3, 1] = calibur.normalized(y) 107 | self.t[:3, 2] = calibur.normalized(z) 108 | return self 109 | 110 | @classmethod 111 | def from_orbit(cls, h, w, radius, azim, elev, origin, fov=30, near=0.1, far=10.0): 112 | """ 113 | Create a perspective camera from an orbital camera. 114 | 115 | Args: 116 | h (int): Height. 117 | w (int): Width. 118 | radius (float): Distance to focus origin. 119 | azim (float): Azimuth angle in degrees relative to focus origin. +Z is 0 degree. 120 | elev (float): Elevation angle in degrees relative to focus origin. 121 | -90 is down-up and 90 is up-down. 122 | Note that -90 and 90 are extreme values where azimuth is undefined, 123 | so use values like 89.999 in these cases instead. 124 | origin (List[float]): List of 3 floats, focus point of the camera. 125 | fov (float): Vertical Field of View in degrees. 126 | near (float): Camera near plane distance. 127 | far (float): Camera far plane distance. 128 | """ 129 | r = radius 130 | cam = cls(h=h, w=w, fov=fov, near=near, far=far) 131 | theta, phi = math.radians(azim), math.radians(elev) 132 | z = r * math.cos(theta) * math.cos(phi) + origin[2] 133 | x = r * math.sin(theta) * math.cos(phi) + origin[0] 134 | y = r * math.sin(phi) + origin[1] 135 | cam.t = trimesh.transformations.translation_matrix([x, y, z]) 136 | cam.lookat(origin) 137 | return cam 138 | -------------------------------------------------------------------------------- /diffrp/rendering/denoiser.py: -------------------------------------------------------------------------------- 1 | """ 2 | Experimental. Implementation and interface subject to change. 3 | """ 4 | ## Copied and modified from intel OIDN 5 | ## Copyright 2018 Intel Corporation 6 | ## SPDX-License-Identifier: Apache-2.0 7 | 8 | import math 9 | import torch 10 | import torch.nn as nn 11 | import torch.nn.functional as F 12 | from ..resources import get_resource_path 13 | from ..utils import * 14 | 15 | 16 | def get_denoiser(): 17 | net = UNet(9, 3) 18 | net.load_state_dict(torch_load_no_warning(get_resource_path("denoisers/rt_hdr_alb_nrm.pt"))) 19 | if torch.cuda.is_available(): 20 | net = net.cuda() 21 | return torch.jit.script(net.float()) 22 | 23 | 24 | def run_denoiser( 25 | denoiser, pbr_hdr: torch.Tensor, albedo_srgb: torch.Tensor, normal: torch.Tensor, 26 | alignment: int = 16 27 | ): 28 | PU_NORM_SCALE = 1. / linear_to_pu(torch.tensor(65504)).item() 29 | in_img = to_bchw(torch.cat([linear_to_pu(pbr_hdr) * PU_NORM_SCALE, albedo_srgb, normal * 0.5 + 0.5], dim=-1)) 30 | b, c, h, w = in_img.shape 31 | dh = math.ceil(h / alignment) * alignment - h 32 | dw = math.ceil(w / alignment) * alignment - w 33 | in_img = nn.ReflectionPad2d((dw // 2, dw - dw // 2, dh // 2, dh - dh // 2))(in_img) 34 | b, c, h, w = in_img.shape 35 | return pu_to_linear(to_hwc(denoiser(in_img))[dh // 2: h - (dh - dh // 2), dw // 2: w - (dw - dw // 2)] / PU_NORM_SCALE) 36 | 37 | 38 | ## ----------------------------------------------------------------------------- 39 | ## Network layers 40 | ## ----------------------------------------------------------------------------- 41 | 42 | 43 | # 3x3 convolution module 44 | def Conv(in_channels, out_channels): 45 | return nn.Conv2d(in_channels, out_channels, 3, padding=1) 46 | 47 | 48 | # ReLU function 49 | def relu(x): 50 | return F.relu(x, inplace=True) 51 | 52 | 53 | # 2x2 max pool function 54 | def pool(x): 55 | return F.max_pool2d(x, 2, 2) 56 | 57 | 58 | # 2x2 nearest-neighbor upsample function 59 | def upsample(x): 60 | return F.interpolate(x, scale_factor=2.0, mode="nearest") 61 | 62 | 63 | # Channel concatenation function 64 | def concat(a, b): 65 | return torch.cat((a, b), 1) 66 | 67 | 68 | ## ----------------------------------------------------------------------------- 69 | ## U-Net model 70 | ## ----------------------------------------------------------------------------- 71 | 72 | 73 | class UNet(nn.Module): 74 | def __init__(self, in_channels, out_channels, small=False): 75 | super(UNet, self).__init__() 76 | 77 | # Number of channels per layer 78 | ic = in_channels 79 | if small: 80 | ec1 = 32 81 | ec2 = 32 82 | ec3 = 32 83 | ec4 = 32 84 | ec5 = 32 85 | dc4 = 64 86 | dc3 = 64 87 | dc2a = 64 88 | dc2b = 32 89 | dc1a = 32 90 | dc1b = 32 91 | else: 92 | ec1 = 32 93 | ec2 = 48 94 | ec3 = 64 95 | ec4 = 80 96 | ec5 = 96 97 | dc4 = 112 98 | dc3 = 96 99 | dc2a = 64 100 | dc2b = 64 101 | dc1a = 64 102 | dc1b = 32 103 | oc = out_channels 104 | 105 | # Convolutions 106 | self.enc_conv0 = Conv(ic, ec1) 107 | self.enc_conv1 = Conv(ec1, ec1) 108 | self.enc_conv2 = Conv(ec1, ec2) 109 | self.enc_conv3 = Conv(ec2, ec3) 110 | self.enc_conv4 = Conv(ec3, ec4) 111 | self.enc_conv5a = Conv(ec4, ec5) 112 | self.enc_conv5b = Conv(ec5, ec5) 113 | self.dec_conv4a = Conv(ec5 + ec3, dc4) 114 | self.dec_conv4b = Conv(dc4, dc4) 115 | self.dec_conv3a = Conv(dc4 + ec2, dc3) 116 | self.dec_conv3b = Conv(dc3, dc3) 117 | self.dec_conv2a = Conv(dc3 + ec1, dc2a) 118 | self.dec_conv2b = Conv(dc2a, dc2b) 119 | self.dec_conv1a = Conv(dc2b + ic, dc1a) 120 | self.dec_conv1b = Conv(dc1a, dc1b) 121 | self.dec_conv0 = Conv(dc1b, oc) 122 | 123 | # Images must be padded to multiples of the alignment 124 | self.alignment = 16 125 | 126 | def forward(self, input): 127 | # Encoder 128 | # ------------------------------------------- 129 | 130 | x = relu(self.enc_conv0(input)) # enc_conv0 131 | 132 | x = relu(self.enc_conv1(x)) # enc_conv1 133 | x = pool1 = pool(x) # pool1 134 | 135 | x = relu(self.enc_conv2(x)) # enc_conv2 136 | x = pool2 = pool(x) # pool2 137 | 138 | x = relu(self.enc_conv3(x)) # enc_conv3 139 | x = pool3 = pool(x) # pool3 140 | 141 | x = relu(self.enc_conv4(x)) # enc_conv4 142 | x = pool(x) # pool4 143 | 144 | # Bottleneck 145 | x = relu(self.enc_conv5a(x)) # enc_conv5a 146 | x = relu(self.enc_conv5b(x)) # enc_conv5b 147 | 148 | # Decoder 149 | # ------------------------------------------- 150 | 151 | x = upsample(x) # upsample4 152 | x = concat(x, pool3) # concat4 153 | x = relu(self.dec_conv4a(x)) # dec_conv4a 154 | x = relu(self.dec_conv4b(x)) # dec_conv4b 155 | 156 | x = upsample(x) # upsample3 157 | x = concat(x, pool2) # concat3 158 | x = relu(self.dec_conv3a(x)) # dec_conv3a 159 | x = relu(self.dec_conv3b(x)) # dec_conv3b 160 | 161 | x = upsample(x) # upsample2 162 | x = concat(x, pool1) # concat2 163 | x = relu(self.dec_conv2a(x)) # dec_conv2a 164 | x = relu(self.dec_conv2b(x)) # dec_conv2b 165 | 166 | x = upsample(x) # upsample1 167 | x = concat(x, input) # concat1 168 | x = relu(self.dec_conv1a(x)) # dec_conv1a 169 | x = relu(self.dec_conv1b(x)) # dec_conv1b 170 | 171 | x = self.dec_conv0(x) # dec_conv0 172 | 173 | return x 174 | 175 | 176 | ## ----------------------------------------------------------------------------- 177 | ## U-Net model: large 178 | ## ----------------------------------------------------------------------------- 179 | 180 | 181 | class UNetLarge(nn.Module): 182 | def __init__(self, in_channels, out_channels, xl=False): 183 | super(UNetLarge, self).__init__() 184 | 185 | # Number of channels per layer 186 | ic = in_channels 187 | if xl: 188 | ec1 = 96 189 | ec2 = 128 190 | ec3 = 192 191 | ec4 = 256 192 | ec5 = 384 193 | dc4 = 256 194 | dc3 = 192 195 | dc2 = 128 196 | dc1 = 96 197 | else: 198 | ec1 = 64 199 | ec2 = 96 200 | ec3 = 128 201 | ec4 = 192 202 | ec5 = 256 203 | dc4 = 192 204 | dc3 = 128 205 | dc2 = 96 206 | dc1 = 64 207 | oc = out_channels 208 | 209 | # Convolutions 210 | self.enc_conv1a = Conv(ic, ec1) 211 | self.enc_conv1b = Conv(ec1, ec1) 212 | self.enc_conv2a = Conv(ec1, ec2) 213 | self.enc_conv2b = Conv(ec2, ec2) 214 | self.enc_conv3a = Conv(ec2, ec3) 215 | self.enc_conv3b = Conv(ec3, ec3) 216 | self.enc_conv4a = Conv(ec3, ec4) 217 | self.enc_conv4b = Conv(ec4, ec4) 218 | self.enc_conv5a = Conv(ec4, ec5) 219 | self.enc_conv5b = Conv(ec5, ec5) 220 | self.dec_conv4a = Conv(ec5 + ec3, dc4) 221 | self.dec_conv4b = Conv(dc4, dc4) 222 | self.dec_conv3a = Conv(dc4 + ec2, dc3) 223 | self.dec_conv3b = Conv(dc3, dc3) 224 | self.dec_conv2a = Conv(dc3 + ec1, dc2) 225 | self.dec_conv2b = Conv(dc2, dc2) 226 | self.dec_conv1a = Conv(dc2 + ic, dc1) 227 | self.dec_conv1b = Conv(dc1, dc1) 228 | self.dec_conv1c = Conv(dc1, oc) 229 | 230 | # Images must be padded to multiples of the alignment 231 | self.alignment = 16 232 | 233 | def forward(self, input): 234 | # Encoder 235 | # ------------------------------------------- 236 | 237 | x = relu(self.enc_conv1a(input)) # enc_conv1a 238 | x = relu(self.enc_conv1b(x)) # enc_conv1b 239 | x = pool1 = pool(x) # pool1 240 | 241 | x = relu(self.enc_conv2a(x)) # enc_conv2a 242 | x = relu(self.enc_conv2b(x)) # enc_conv2b 243 | x = pool2 = pool(x) # pool2 244 | 245 | x = relu(self.enc_conv3a(x)) # enc_conv3a 246 | x = relu(self.enc_conv3b(x)) # enc_conv3b 247 | x = pool3 = pool(x) # pool3 248 | 249 | x = relu(self.enc_conv4a(x)) # enc_conv4a 250 | x = relu(self.enc_conv4b(x)) # enc_conv4b 251 | x = pool(x) # pool4 252 | 253 | # Bottleneck 254 | x = relu(self.enc_conv5a(x)) # enc_conv5a 255 | x = relu(self.enc_conv5b(x)) # enc_conv5b 256 | 257 | # Decoder 258 | # ------------------------------------------- 259 | 260 | x = upsample(x) # upsample4 261 | x = concat(x, pool3) # concat4 262 | x = relu(self.dec_conv4a(x)) # dec_conv4a 263 | x = relu(self.dec_conv4b(x)) # dec_conv4b 264 | 265 | x = upsample(x) # upsample3 266 | x = concat(x, pool2) # concat3 267 | x = relu(self.dec_conv3a(x)) # dec_conv3a 268 | x = relu(self.dec_conv3b(x)) # dec_conv3b 269 | 270 | x = upsample(x) # upsample2 271 | x = concat(x, pool1) # concat2 272 | x = relu(self.dec_conv2a(x)) # dec_conv2a 273 | x = relu(self.dec_conv2b(x)) # dec_conv2b 274 | 275 | x = upsample(x) # upsample1 276 | x = concat(x, input) # concat1 277 | x = relu(self.dec_conv1a(x)) # dec_conv1a 278 | x = relu(self.dec_conv1b(x)) # dec_conv1b 279 | x = relu(self.dec_conv1c(x)) # dec_conv1c 280 | 281 | return x 282 | -------------------------------------------------------------------------------- /diffrp/rendering/interpolator.py: -------------------------------------------------------------------------------- 1 | """ 2 | Implementation of different varying variable interpolators. 3 | 4 | This is an implementation detail and has no guarantee on API stability. 5 | 6 | See ``interpolator_impl`` in :py:class:`diffrp.rendering.surface_deferred.SurfaceDeferredRenderSessionOptions` 7 | for explanation on selecting interpolators. 8 | """ 9 | import abc 10 | import torch 11 | from typing import Optional 12 | import nvdiffrast.torch as dr 13 | from ..utils import fma 14 | 15 | 16 | def __float_as_int(x: torch.Tensor): 17 | return x.view(torch.int32) 18 | 19 | 20 | def __int_as_float(x: torch.Tensor): 21 | return x.view(torch.float32) 22 | 23 | 24 | def float_to_triidx(x: torch.Tensor): 25 | return torch.where(x <= 16777216, x.int(), __float_as_int(x) - 0x4a800000) 26 | 27 | 28 | def triidx_to_float(x: torch.Tensor): 29 | return torch.where(x <= 0x01000000, x.float(), __int_as_float(0x4a800000 + x)) 30 | 31 | 32 | @torch.jit.script 33 | def _interpolate_impl( 34 | vertex_buffer: torch.Tensor, 35 | vi_data: torch.Tensor, 36 | tris_idx: torch.Tensor, 37 | empty_region: Optional[float] = None, 38 | ): 39 | v123 = vertex_buffer[tris_idx] 40 | v1, v2, v3 = torch.select(v123, -2, 0), torch.select(v123, -2, 1), torch.select(v123, -2, 2) 41 | # v1: *, d 42 | u, v = vi_data[..., 0:1], vi_data[..., 1:2] 43 | # result = v1 * u + v2 * v + v3 * (1 - u - v) 44 | # (v1 - v3) * u + (v2 - v3) * v + v3 45 | result = fma(v1 - v3, u, fma(v2 - v3, v, v3)) 46 | if empty_region is not None: 47 | result = torch.where(vi_data[..., -1:] > 0, result, empty_region) 48 | return result 49 | 50 | 51 | def polyfill_interpolate( 52 | vertex_buffer: torch.Tensor, 53 | vi_data: torch.Tensor, 54 | tri_idx: torch.IntTensor, 55 | empty_region: Optional[float] = None, 56 | ): 57 | # n, d; f, 3; *, 4 -> *, d 58 | # tris[vi]: *, 3 59 | # v[tris[vi]]: *, 3, d 60 | # v1, v2, v3 = torch.unbind(vertex_buffer[tris[vi_idx - 1]], dim=-2) 61 | # v1: *, d 62 | # result = v1 * vi_data.x + v2 * vi_data.y + v3 * (1 - vi_data.x - vi_data.y) 63 | # if empty_region is not None: 64 | # result = torch.where(vi_data.w > 0, result, empty_region) 65 | return _interpolate_impl(vertex_buffer, vi_data, tri_idx, empty_region) 66 | 67 | 68 | class Interpolator(metaclass=abc.ABCMeta): 69 | vi_data: torch.Tensor 70 | 71 | @abc.abstractmethod 72 | def interpolate(self, vertex_buffer: torch.Tensor) -> torch.Tensor: 73 | raise NotImplementedError 74 | 75 | 76 | class FullScreenInterpolator(Interpolator): 77 | 78 | def __init__(self, vi_data: torch.Tensor, tris: torch.IntTensor): 79 | self.vi_data = vi_data.contiguous() # (u, v, z|depth, tridx) 80 | self.tris = tris.contiguous() 81 | 82 | def interpolate(self, vertex_buffer: torch.Tensor): 83 | attr, da = dr.interpolate(vertex_buffer.contiguous(), self.vi_data, self.tris) 84 | return attr 85 | 86 | 87 | class MaskedSparseInterpolator(Interpolator): 88 | 89 | def __init__(self, vi_data: torch.Tensor, tris: torch.IntTensor, mask: torch.BoolTensor): 90 | self.tris = tris 91 | self.indices = mask.nonzero(as_tuple=True) 92 | self.vi_data = vi_data[self.indices] # (u, v, z|depth, tridx) 93 | self.tri_idx = tris[float_to_triidx(self.vi_data[..., -1]) - 1].long() 94 | 95 | def interpolate(self, vertex_buffer: torch.Tensor): 96 | # n, d; f, 3; ?, 4 97 | return polyfill_interpolate(vertex_buffer, self.vi_data, self.tri_idx) 98 | -------------------------------------------------------------------------------- /diffrp/rendering/mixin.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from ..scene import Scene 3 | from .camera import Camera 4 | from ..utils.cache import cached 5 | from ..utils.shader_ops import * 6 | from typing import List, Callable, Union, Tuple 7 | from .interpolator import MaskedSparseInterpolator 8 | from ..utils.coordinates import near_plane_ndc_grid 9 | from ..materials.base_material import VertexArrayObject, SurfaceInput, SurfaceOutputStandard 10 | 11 | 12 | class RenderSessionMixin: 13 | """ 14 | Mixin implementation for common operations across render sessions. 15 | """ 16 | scene: Scene 17 | camera: Camera 18 | 19 | @cached 20 | def camera_VP(self): 21 | return torch.mm(self.camera_P(), self.camera_V()) 22 | 23 | @cached 24 | def camera_P(self): 25 | return self.camera.P() 26 | 27 | @cached 28 | def camera_V(self): 29 | return self.camera.V() 30 | 31 | @staticmethod 32 | def _view_dir_impl(grid: torch.Tensor, v: torch.Tensor, p: torch.Tensor, vp: torch.Tensor): 33 | inved_matrices = small_matrix_inverse(torch.stack([v, vp])) 34 | camera_pos = inved_matrices[0, :3, 3] 35 | grid = torch.matmul(grid, inved_matrices[1].T) 36 | grid = grid[..., :3] / grid[..., 3:] 37 | rays_d = normalized(grid - camera_pos) 38 | near = p[2, 3] / (p[2, 2] - 1) 39 | return camera_pos + rays_d * near, rays_d 40 | 41 | @cached 42 | def camera_far(self) -> float: 43 | p = self.camera_P() 44 | return (p[2, 3] / (p[2, 2] + 1)).item() 45 | 46 | @cached 47 | def camera_rays(self) -> torch.Tensor: 48 | grid = near_plane_ndc_grid(*self.camera.resolution(), torch.float32, torch.device('cuda')) 49 | return self._view_dir_impl(grid, self.camera_V(), self.camera_P(), self.camera_VP()) 50 | 51 | @cached 52 | def view_dir(self) -> torch.Tensor: 53 | """ 54 | Returns: 55 | torch.Tensor: 56 | Full-screen unit direction vectors of outbound rays from the camera. 57 | Tensor of shape (H, W, 3). 58 | """ 59 | return self.camera_rays()[1] 60 | 61 | def _cat_fast_path(self, attrs: List[torch.Tensor]): 62 | if len(attrs) == 1: 63 | return attrs[0] 64 | else: 65 | return torch.cat(attrs) 66 | 67 | def _checked_cat(self, attrs: List[torch.Tensor], verts_ref: List[torch.Tensor], expected_dim): 68 | for v, a in zip(verts_ref, attrs): 69 | assert v.shape[0] == a.shape[0], "attribute length not the same as number of vertices" 70 | if isinstance(expected_dim, int): 71 | assert a.shape[-1] == expected_dim, "expected %d dims but got %d for vertex attribute" % (expected_dim, a.shape[-1]) 72 | return self._cat_fast_path(attrs) 73 | 74 | @cached 75 | def vertex_array_object(self) -> VertexArrayObject: 76 | world_pos = [] 77 | tris = [] 78 | sts = [torch.full([1], 0, device='cuda', dtype=torch.int64)] 79 | objs = self.scene.objects 80 | offset = 0 81 | for s, x in enumerate(objs): 82 | world_pos.append(transform_point4x3(x.verts, x.M)) 83 | assert x.tris.shape[-1] == 3, "Expected 3 vertices per triangle, got %d" % x.tris.shape[-1] 84 | tris.append(x.tris + offset) 85 | sts.append(torch.full([len(x.tris)], s + 1, dtype=torch.int64, device=x.tris.device)) 86 | offset += len(x.verts) 87 | total_custom_attrs = set().union(*(obj.custom_attrs.keys() for obj in objs)) 88 | for k in total_custom_attrs: 89 | for obj in objs: 90 | if k in obj.custom_attrs: 91 | sz = obj.custom_attrs[k].shape[-1] 92 | break 93 | else: 94 | assert False, "Internal assertion failure. You should never reach here." 95 | for obj in objs: 96 | if k in obj.custom_attrs: 97 | assert len(obj.custom_attrs[k]) == len(obj.verts), "Attribute length not the same as number of vertices: %s" % k 98 | assert obj.custom_attrs[k].shape[-1] == sz, "Different dimensions for same custom attribute: %d != %d" % (obj.custom_attrs[k].shape[-1], sz) 99 | else: 100 | obj.custom_attrs[k] = zeros_like_vec(obj.verts, sz) 101 | 102 | verts_ref = [obj.verts for obj in objs] 103 | return VertexArrayObject( 104 | self._checked_cat([obj.verts for obj in objs], verts_ref, 3), 105 | self._checked_cat([obj.normals for obj in objs], verts_ref, 3), 106 | self._cat_fast_path(world_pos), 107 | self._cat_fast_path(tris).int(memory_format=torch.contiguous_format), 108 | self._cat_fast_path(sts).int(memory_format=torch.contiguous_format), 109 | self._checked_cat([obj.color for obj in objs], verts_ref, None), 110 | self._checked_cat([obj.uv for obj in objs], verts_ref, 2), 111 | self._checked_cat([obj.tangents for obj in objs], verts_ref, 4), 112 | {k: torch.cat([obj.custom_attrs[k] for obj in objs]) for k in total_custom_attrs} 113 | ) 114 | 115 | def _collector_world_normal(self, si: SurfaceInput, so: SurfaceOutputStandard): 116 | if so.normal is None: 117 | return si.world_normal 118 | if so.normal_space == 'tangent': 119 | vn = si.world_normal_unnormalized 120 | vnt = so.normal 121 | vt, vs = si.world_tangent.xyz, si.world_tangent.w 122 | vb = vs * cross(vn, vt) 123 | return normalized(fma(vnt.x, vt, fma(vnt.y, vb, vnt.z * vn))) 124 | if so.normal_space == 'object': 125 | return normalized(transform_vector3x3(so.normal, si.uniforms.M)) 126 | if so.normal_space == 'world': 127 | return normalized(so.normal) 128 | assert False, "Unknown normal space: " + so.normal_space 129 | 130 | def _gbuffer_collect_layer_impl_stencil_masked( 131 | self, 132 | mats: List[Tuple[SurfaceInput, SurfaceOutputStandard]], 133 | operator: Callable[[SurfaceInput, SurfaceOutputStandard], torch.Tensor], 134 | initial_buffer: torch.Tensor 135 | ): 136 | buffer = initial_buffer 137 | indices = [] 138 | values = [] 139 | nind = 0 140 | for si, so in mats: 141 | interpolator: MaskedSparseInterpolator = si.interpolator 142 | if len(si.interpolator.indices[0]) == 0: 143 | continue 144 | op = operator(si, so) 145 | if op is not None: 146 | nind = len(interpolator.indices) 147 | indices.append(interpolator.indices) 148 | values.append(op) 149 | if len(indices) == 1: 150 | buffer = buffer.index_put(indices[0], values[0]) 151 | elif nind != 0: 152 | indices = tuple(torch.cat([x[i] for x in indices]) for i in range(nind)) 153 | values = torch.cat(values) 154 | buffer = buffer.index_put(indices, values) 155 | return buffer 156 | -------------------------------------------------------------------------------- /diffrp/rendering/path_tracing.py: -------------------------------------------------------------------------------- 1 | """ 2 | Experimental. Implementation and interface subject to change. 3 | """ 4 | import math 5 | import torch 6 | from dataclasses import dataclass 7 | from typing_extensions import Literal 8 | from typing import Optional, Dict, Callable 9 | 10 | from .camera import Camera 11 | from ..utils.cache import cached 12 | from ..utils.shader_ops import * 13 | from .mixin import RenderSessionMixin 14 | from ..utils.geometry import barycentric 15 | from ..scene import Scene, ImageEnvironmentLight 16 | from ..utils.raycaster import NaivePBBVH, BruteForceRaycaster, TorchOptiX 17 | from .interpolator import MaskedSparseInterpolator, triidx_to_float, float_to_triidx 18 | from ..utils.coordinates import near_plane_ndc_grid, unit_direction_to_latlong_uv 19 | from ..materials.base_material import SurfaceInput, SurfaceUniform, SurfaceOutputStandard 20 | from ..utils.light_transport import hammersley, importance_sample_ggx, combine_fixed_tangent_space, geometry_smith, fresnel_schlick_smoothness 21 | 22 | 23 | @dataclass 24 | class PathTracingSessionOptions: 25 | """ 26 | Options for ``PathTracingSession``. 27 | 28 | Args: 29 | ray_depth (int): 30 | The tracing depth of a ray. 31 | Defaults to 3, meaning at maximum 3 bounces of a ray will be computed. 32 | At the end of the bounce, ``pbr_ray_last_bounce`` determines the behavior of the ray if you are using the default ``sampler_brdf``. 33 | Otherwise, the behavior is defined by your implementation of the sampler. 34 | ray_spp (int): 35 | Samples per pixel. Defaults to 16. 36 | ray_split_size (int): 37 | Split the rays traced as a batch. 38 | Useful to reduce memory consumption when gradients are not required. 39 | Defaults to 8M (8 x 1024 x 1024) rays. 40 | deterministic (bool): 41 | Try to make the render deterministic. 42 | Defaults to True. 43 | pbr_ray_step_epsilon (float): 44 | To avoid near-intersection at the origin (surface bounce point) of bounce rays, we march ray origins by this amount when generating them in ``sampler_brdf``. 45 | You are also recommended to reuse this value for your own sampler for bounce rays. 46 | Defaults to 1e-3 (0.001). 47 | pbr_ray_last_bounce (str): 48 | One of 'void' and 'skybox'. 49 | Special to ``sampler_brdf`` in PBR rendering. 50 | If set to 'void' (default), a bounce ray ends with zero radiance if it has not reached the sky in ``ray_depth`` bounces. 51 | If set to 'skybox', it ends with the sky radiance in the final-bounce direction instead. 52 | raycaster_impl (str): 53 | | One of 'brute-force', 'naive-pbbvh' and 'torchoptix'. 54 | Defaults to 'torchoptix'. 55 | | The implementation for scene raycaster. 56 | | 'torchoptix' requires ``torchoptix`` to be installed. It uses hardware RT cores for raycasting. This is the fastest implementation. 57 | | 'naive-pbbvh' uses a software implementation of a GPU BVH. It is usually 10x to 100x slower than the hardware solution of 'torchoptix'. 58 | | 'brute-force' uses a software brute-force ray-triangle tests comparing all pairs of rays and triangles. It is fast if there is little geometry (fewer than ~10 triangles). 59 | raycaster_epsilon (float): 60 | Only relevant for 'brute-force' and 'naive-pbbvh' raycaster implementations. Discards triangles parallel to the rays beyond this threshold. 61 | Defaults to 1e-8. 62 | raycaster_builder (str): 63 | | One of 'morton' and 'splitaxis'. 64 | Defaults to 'splitaxis'. 65 | | The BVH builder. Only relevant for 'naive-pbbvh' raycaster implementation. 66 | | 'morton' uses a linear BVH sorting triangles based on 30-bit morton code. This has lower query performance but faster building. 67 | | 'splitaxis' sorts and splits triangles based on the longest axis at each BVH level. 68 | optix_log_level (int): 69 | OptiX log level, 0 to 4. Higher means more verbose. 70 | Defaults to 3. 71 | """ 72 | ray_depth: int = 3 73 | ray_spp: int = 16 74 | ray_split_size: int = 8 * 1024 * 1024 75 | 76 | deterministic: bool = True 77 | 78 | pbr_ray_step_epsilon: float = 1e-3 79 | pbr_ray_last_bounce: Literal['void', 'skybox'] = 'void' 80 | 81 | raycaster_impl: Literal['brute-force', 'naive-pbbvh', 'torchoptix'] = 'torchoptix' 82 | raycaster_epsilon: float = 1e-8 83 | raycaster_builder: Literal['morton', 'splitaxis'] = 'splitaxis' 84 | optix_log_level: Literal[0, 1, 2, 3, 4] = 3 85 | 86 | 87 | @dataclass 88 | class RayOutputs: 89 | """ 90 | The protocol outputs for path tracing sampler. 91 | 92 | Args: 93 | radiance (torch.Tensor): 94 | Shape (R, C) where R is number of rays for the current trace and C is specified channels in ``trace_rays``. 95 | Radiance value of the current hit. 96 | The output radiance would be the sum of radiance times the attenuation at the radiance. 97 | transfer (torch.Tensor): 98 | Should be brocastable with radiance. 99 | Transfer value of the current hit. 100 | The total attenuation value would be the product of transfer values. 101 | The attenuation value for the current radiance excludes the current transfer function (starts with all-1s for the first-hit radiance). 102 | next_rays_o (torch.Tensor): 103 | Shape (R, 3), bounce ray origins. 104 | You may want to advance ray origins by a small value (e.g., ``pbr_ray_step_epsilon``) to avoid intersecting the same point at the origin for secondary bounces. 105 | next_rays_d (torch.Tensor): 106 | Shape (R, 3), bounce ray directions. 107 | alpha (torch.Tensor): 108 | Shape (R, 1), alpha values. 109 | Additively composed and saturated for a hit-mask output. 110 | Has no exact physical meaning for different values between 0 and 1 in path tracing, 111 | and has no influence on radiance composition. 112 | extras (Dict[str, torch.Tensor]): 113 | Extra auxiliary outputs indexed by a string name. 114 | Collected on the first hit and ignoring . 115 | Useful for auxiliary outputs like normals, base colors and AOVs. 116 | See also: 117 | :py:meth:`diffrp.rendering.path_tracing.PathTracingSession.trace_rays` 118 | """ 119 | radiance: torch.Tensor 120 | transfer: torch.Tensor 121 | next_rays_o: torch.Tensor 122 | next_rays_d: torch.Tensor 123 | alpha: torch.Tensor 124 | extras: Dict[str, torch.Tensor] 125 | 126 | 127 | class PathTracingSession(RenderSessionMixin): 128 | """ 129 | Path tracing render pipeline in DiffRP. 130 | """ 131 | def __init__( 132 | self, 133 | scene: Scene, camera: Camera, 134 | options: Optional[PathTracingSessionOptions] = None 135 | ) -> None: 136 | self.scene = scene 137 | self.camera = camera 138 | if options is None: 139 | options = PathTracingSessionOptions() 140 | self.options = options 141 | 142 | @cached 143 | def raycaster(self): 144 | vao = self.vertex_array_object() 145 | cfg = {'epsilon': self.options.raycaster_epsilon} 146 | if self.options.raycaster_impl == 'torchoptix': 147 | cfg['optix_log_level'] = self.options.optix_log_level 148 | return TorchOptiX(vao.world_pos, vao.tris, cfg) 149 | elif self.options.raycaster_impl == 'naive-pbbvh': 150 | if len(vao.tris) <= 3: 151 | return BruteForceRaycaster(vao.world_pos, vao.tris, cfg) 152 | else: 153 | cfg['builder'] = self.options.raycaster_builder 154 | return NaivePBBVH(vao.world_pos, vao.tris, cfg) 155 | else: 156 | return BruteForceRaycaster(vao.world_pos, vao.tris, cfg) 157 | 158 | def layer_material_rays(self, rays_o, rays_d, t, i: torch.Tensor): 159 | cam_v = self.camera_V() 160 | cam_p = self.camera_P() 161 | far = self.camera_far() 162 | vao = self.vertex_array_object() 163 | i = torch.where(t < far, i + 1, 0) 164 | vi_data = float4( 165 | barycentric(*torch.unbind(vao.world_pos[vao.tris[i - 1]], dim=-2), rays_o + rays_d * t[..., None]).xy, 166 | t[..., None], triidx_to_float(i[..., None]) 167 | ).view(-1, 4) 168 | mats = [] 169 | stencil_buf = vao.stencils[i] 170 | for pi, x in enumerate(self.scene.objects): 171 | su = SurfaceUniform(x.M, cam_v, cam_p) 172 | si = SurfaceInput(su, vao, MaskedSparseInterpolator(vi_data, vao.tris, stencil_buf == pi + 1)) 173 | if len(si.interpolator.indices[0]) == 0: 174 | mats.append((si, SurfaceOutputStandard())) 175 | continue 176 | so = x.material.shade(su, si) 177 | mats.append((si, so)) 178 | return mats 179 | 180 | def _super_collector(self, si: SurfaceInput, so: SurfaceOutputStandard): 181 | albedo = so.albedo if so.albedo is not None else gpu_f32([1.0, 0.0, 1.0]).expand_as(world_normal) 182 | world_normal = self._collector_world_normal(si, so) 183 | metal = so.metallic if so.metallic is not None else zeros_like_vec(world_normal, 1) 184 | smooth = so.smoothness if so.smoothness is not None else full_like_vec(world_normal, 0.5, 1) 185 | alpha = so.alpha if so.alpha is not None else full_like_vec(world_normal, 1, 1) 186 | emission = so.emission if so.emission is not None else full_like_vec(world_normal, 0, 3) 187 | return torch.cat([albedo, world_normal, metal, smooth, alpha, emission], dim=-1) 188 | 189 | @staticmethod 190 | @torch.jit.script 191 | def _sampler_brdf_impl(attrs: torch.Tensor, t: torch.Tensor, rays_o: torch.Tensor, rays_d: torch.Tensor, env_radiance: torch.Tensor): 192 | albedo = attrs[..., 0:3] 193 | world_normal = attrs[..., 3:6] 194 | metal = attrs[..., 6:7] 195 | smooth = attrs[..., 7:8] 196 | alpha = attrs[..., 8:9] 197 | emission = attrs[..., 9:12] 198 | next_o = rays_o + rays_d * t[..., None] 199 | 200 | dielectric = 1 - metal 201 | diffuse_color = dielectric * albedo 202 | d = diffuse_color.max(-1, True).values 203 | diffuse_prob = dielectric * d / (0.04 + d) 204 | specular_prob = 1 - diffuse_prob # 0.04 * dielectric + metal 205 | ray_is_transmit = torch.rand_like(alpha) >= alpha 206 | ray_is_diffuse = torch.rand_like(metal) >= specular_prob 207 | 208 | # transmit 209 | next_d_tr = rays_d 210 | transfer_tr = 1.0 211 | 212 | # diffuse 213 | z2 = torch.rand_like(metal) 214 | theta = torch.rand_like(z2) * math.tau 215 | xy = torch.sqrt(1.0 - z2) 216 | next_d_di = combine_fixed_tangent_space(xy * torch.cos(theta), z2.sqrt(), xy * torch.sin(theta), world_normal) 217 | diffuse_color = albedo * dielectric 218 | transfer_di = diffuse_color / torch.clamp_min(diffuse_prob, 0.0001) 219 | 220 | # specular 221 | roughness_clip = (1 - smooth).clamp_min(1 / 512) 222 | x = torch.rand_like(smooth) 223 | y = torch.rand_like(smooth) 224 | h = importance_sample_ggx(x, y, world_normal, roughness_clip) 225 | next_d_sp = reflect(rays_d, h) 226 | dot_v_h = -dot(h, rays_d) 227 | f0 = fma(albedo, metal, 0.04 * dielectric) 228 | transfer_sp = ( 229 | fresnel_schlick_smoothness(dot_v_h, f0, smooth) * 230 | geometry_smith(world_normal, -rays_d, next_d_sp, roughness_clip) * 231 | torch.clamp_min(dot_v_h, 1e-6) / (torch.clamp_min(dot(world_normal, h), 1e-6) * torch.clamp_min(-dot(world_normal, rays_d), 1e-6)) 232 | ) / torch.clamp_min(specular_prob * alpha, 0.0001) 233 | 234 | next_d = torch.where(ray_is_transmit, next_d_tr, torch.where(ray_is_diffuse, next_d_di, next_d_sp)) 235 | transfer = torch.where(ray_is_transmit, transfer_tr, torch.where(ray_is_diffuse, transfer_di, transfer_sp)) 236 | return albedo, emission, world_normal, alpha, emission + env_radiance, transfer, next_o, next_d 237 | 238 | @cached 239 | def _single_env_light(self): 240 | env_light = None 241 | for light in self.scene.lights: 242 | if isinstance(light, ImageEnvironmentLight): 243 | if env_light is not None: 244 | raise ValueError("Only one environment light is supported in path tracing now.") 245 | env_light = light.image_rh() 246 | if env_light is None: 247 | return black_tex().rgb 248 | return env_light 249 | 250 | def sampler_brdf(self, rays_o, rays_d, t, i, d: int) -> RayOutputs: 251 | """ 252 | A *sampler* for ``trace_rays`` implementing the principled PBR BRDF. 253 | 254 | See :py:class:`RayOutputs` and :py:meth:`~diffrp.rendering.path_tracing.PathTracingSession.pbr` for semantics of outputs. 255 | 256 | See also: 257 | :py:meth:`~diffrp.rendering.path_tracing.PathTracingSession.trace_rays` for meaning of arguments. 258 | """ 259 | far = self.camera_far() 260 | mats = self.layer_material_rays(rays_o, rays_d, t, i) 261 | always_sky = self.options.ray_depth - 1 == d and self.options.pbr_ray_last_bounce == 'skybox' 262 | attrs = self._gbuffer_collect_layer_impl_stencil_masked( 263 | mats, self._super_collector, zeros_like_vec(rays_o, 12) 264 | ) 265 | 266 | hit = t[..., None] < far 267 | env_radiance = sample2d(self._single_env_light(), float2(*unit_direction_to_latlong_uv(rays_d))) 268 | if not always_sky: 269 | env_radiance = torch.where(hit, 0.0, env_radiance) 270 | albedo, emission, world_normal, alpha, radiance, transfer, next_o, next_d = PathTracingSession._sampler_brdf_impl(attrs, t, rays_o, rays_d, env_radiance) 271 | 272 | return RayOutputs( 273 | alpha=alpha, 274 | radiance=radiance, 275 | transfer=torch.where(hit, transfer, 0.0), 276 | next_rays_o=next_o + next_d * self.options.pbr_ray_step_epsilon, 277 | next_rays_d=next_d, 278 | extras=dict(albedo=albedo, emission=emission, world_normal=world_normal, world_position=next_o) 279 | ) 280 | 281 | def trace_rays(self, sampler: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int], RayOutputs], radiance_channels: int = 3): 282 | """ 283 | Trace the ray paths. 284 | 285 | Note that missed rays are not automatically excluded in further traces. 286 | You may need to specify a transfer value of zeros to avoid light leakage. 287 | 288 | Args: 289 | sampler (Callable): 290 | | Sampler function to produce radiance, alpha and extra values. 291 | Take five arguments ``(rays_o, rays_d, t, i, d)`` and outputs :py:class:`RayOutputs`. 292 | | ``rays_o`` and ``rays_d`` are of shape (R, 3), and defines current tracing rays. 293 | | ``t`` of shape (R) is the distance from the ray origin to the hit point. It is guaranteed to be ``camera_far()`` if not hit. 294 | | ``i`` of shape (R) is the primitive index of the hit point, starting from 0. The value is undefined for non-hit rays. 295 | | Usually you do not need to handle the semantics of ``i`` yourselves, but 296 | use ``layer_material_rays`` and ``_gbuffer_collect_layer_impl_stencil_masked`` to collect information about the scene. 297 | | ``d`` is the current ray depth. ``0`` is the first rays emitting from the camera. 298 | | Returns :py:class:`RayOutputs`. See the documentation for :py:class:`RayOutputs` for meanings of components. 299 | radiance_channels (int): 300 | Number of channels for radiance outputs in sampler. 301 | Defaults to 3. 302 | 303 | Returns: 304 | Tuple of (radiance, alpha, extras): 305 | ``radiance`` is a tensor of shape (H, W, ``radiance_channels``). 306 | ``alpha`` is a tensor of shape (H, W, 1). 307 | ``extras`` is a dict from names to tensors of auxiliary attributes, 308 | all of which has shape (H, W, C) where C is the number of channels output by the ``sampler``. 309 | """ 310 | device = torch.device('cuda') 311 | raycaster = self.raycaster() 312 | far = self.camera_far() 313 | cam_v = self.camera_V() 314 | cam_p = self.camera_P() 315 | H, W = self.camera.resolution() 316 | pix_grid = near_plane_ndc_grid(H, W, torch.float32, device).reshape(-1, 4) 317 | qrx, qry = hammersley(self.options.ray_spp, self.options.deterministic, device) 318 | sections = min(self.options.ray_spp, math.ceil((H * W * self.options.ray_spp) / self.options.ray_split_size)) 319 | qrx = torch.tensor_split(qrx, sections) 320 | qry = torch.tensor_split(qry, sections) 321 | radiance = zeros_like_vec(pix_grid, radiance_channels) 322 | alpha = zeros_like_vec(pix_grid, 1) 323 | extras_c = dict() 324 | extras_v = dict() 325 | for qrx, qry in zip(qrx, qry): 326 | espp = len(qrx) 327 | qrx, qry = qrx[..., None, None], qry[..., None, None] 328 | grid = pix_grid 329 | grid = float4(grid.x + (qrx - 0.5) * (2 / W), grid.y + (qry - 0.5) * (2 / H), grid.zw) # espp, HW, 4 330 | grid = grid.reshape(-1, 4) 331 | rays_o, rays_d = self._view_dir_impl(grid, cam_v, cam_p, self.camera_VP()) 332 | transfer = ones_like_vec(rays_o, radiance_channels) 333 | for d in range(self.options.ray_depth): 334 | t, i = raycaster.query(rays_o, rays_d, far) 335 | ray_outputs = sampler(rays_o, rays_d, t, i, d) 336 | radiance = radiance + (transfer * ray_outputs.radiance).view(espp, -1, radiance_channels).sum(0) 337 | alpha = alpha + ray_outputs.alpha.reshape(espp, -1, 1).sum(0) 338 | transfer = transfer * ray_outputs.transfer 339 | rays_o, rays_d = ray_outputs.next_rays_o, ray_outputs.next_rays_d 340 | if d == 0: 341 | for k, v in ray_outputs.extras.items(): 342 | if k not in extras_c: 343 | extras_c[k] = espp 344 | extras_v[k] = v.reshape(espp, H, W, v.shape[-1]).sum(0) 345 | else: 346 | extras_c[k] += espp 347 | extras_v[k] = extras_v[k] + v.reshape(espp, H, W, v.shape[-1]).sum(0) 348 | radiance = radiance.reshape(H, W, radiance_channels) / self.options.ray_spp 349 | alpha = saturate(alpha.reshape(H, W, 1) / self.options.ray_spp) 350 | for k in extras_v: 351 | extras_v[k] = torch.flipud(extras_v[k] / extras_c[k]) 352 | return torch.flipud(radiance), torch.flipud(alpha), extras_v 353 | 354 | def pbr(self): 355 | """ 356 | Path-traced PBR rendering. 357 | 358 | Returns: 359 | Tuple of (radiance, alpha, extras): 360 | | ``radiance`` is a linear HDR RGB tensor of shape (H, W, 3). 361 | | ``alpha`` is a mask tensor of shape (H, W, 1). 362 | | ``extras`` is a dict from names to tensors of auxiliary attributes. 363 | The included attributes are: ``albedo``, ``emission``, ``world_normal`` and ``world_position``. 364 | ``world_normal`` are raw vectors in range of [-1, 1] and are not mapped to false colors. 365 | All of them have shape (H, W, 3). 366 | """ 367 | return self.trace_rays(self.sampler_brdf) 368 | -------------------------------------------------------------------------------- /diffrp/resources/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | 4 | def get_resource_path(rel): 5 | """ 6 | :meta private: 7 | """ 8 | return os.path.join(os.path.dirname(os.path.abspath(__file__)), rel) 9 | -------------------------------------------------------------------------------- /diffrp/resources/denoisers/rt_hdr_alb_nrm.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/diffrp/resources/denoisers/rt_hdr_alb_nrm.pt -------------------------------------------------------------------------------- /diffrp/resources/hdri_exrs/newport_loft.exr: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/diffrp/resources/hdri_exrs/newport_loft.exr -------------------------------------------------------------------------------- /diffrp/resources/hdris.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import pyexr 3 | from . import get_resource_path 4 | from ..utils.cache import singleton_cached 5 | 6 | 7 | @singleton_cached 8 | def newport_loft(): 9 | """ 10 | Returns: 11 | torch.Tensor: 12 | a CPU float32 tensor of shape (800, 1600, 3) representing the environment of the NewPort Loft scene. 13 | """ 14 | with pyexr.open(get_resource_path("hdri_exrs/newport_loft.exr")) as fi: 15 | return torch.tensor(fi.get()) 16 | -------------------------------------------------------------------------------- /diffrp/resources/luts/agx-base-contrast.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/diffrp/resources/luts/agx-base-contrast.pt -------------------------------------------------------------------------------- /diffrp/scene/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Includes the scene description components. 3 | The APIs in the subpackages are also directly available from this ``diffrp.scene`` package. 4 | """ 5 | 6 | from .scene import Scene 7 | from .lights import DirectionalLight, PointLight, ImageEnvironmentLight 8 | from .objects import MeshObject 9 | -------------------------------------------------------------------------------- /diffrp/scene/lights.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Union 3 | from dataclasses import dataclass 4 | 5 | 6 | @dataclass 7 | class Light(object): 8 | """ 9 | Base class for Lights. 10 | 11 | Args: 12 | intensity (float): Linear multiplier of light intensity. 13 | color (torch.Tensor): Shape (3,). Linear RGB tint of light. 14 | """ 15 | intensity: Union[torch.Tensor, float] 16 | color: torch.Tensor 17 | 18 | 19 | @dataclass 20 | class DirectionalLight(Light): 21 | """ 22 | Directional light data structure; Not supported yet. 23 | """ 24 | direction: torch.Tensor 25 | 26 | 27 | @dataclass 28 | class PointLight(Light): 29 | """ 30 | Point light data structure; Not supported yet. 31 | """ 32 | position: torch.Tensor 33 | 34 | 35 | @dataclass 36 | class ImageEnvironmentLight(Light): 37 | """ 38 | Environment light based on lat-long format images. 39 | 40 | Args: 41 | intensity (float): Linear multiplier of light intensity. 42 | color (torch.Tensor): Shape (3,). Linear RGB tint of light. 43 | image (torch.Tensor): (H, 2H, 3) RGB linear values of lighting environment. 44 | render_skybox (bool): Whether to show this environment in PBR as the background. 45 | """ 46 | image: torch.Tensor 47 | render_skybox: bool = True 48 | 49 | def image_rh(self): 50 | return torch.fliplr(self.image) * (self.intensity * self.color) 51 | -------------------------------------------------------------------------------- /diffrp/scene/objects.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from dataclasses import dataclass 3 | from typing_extensions import Literal 4 | from typing import Optional, Dict, Union, Any, TYPE_CHECKING 5 | from ..utils.shader_ops import zeros_like_vec, ones_like_vec 6 | from ..utils.geometry import make_face_soup, compute_face_normals, compute_vertex_normals 7 | if TYPE_CHECKING: 8 | from ..materials.base_material import SurfaceMaterial 9 | 10 | 11 | @dataclass 12 | class MeshObject: 13 | """ 14 | Represents a triangle mesh object in the scene. 15 | 16 | Args: 17 | material (SurfaceMaterial): 18 | The material to render with. 19 | You can use the :py:class:`diffrp.materials.default_material.DefaultMaterial` 20 | if you only care about the triangle geometry. 21 | verts (torch.Tensor): 22 | GPU tensor of vertices. Tensor of shape (V, 3), dtype float32. 23 | tris (torch.IntTensor): 24 | GPU tensor of indices. Tensor of shape (F, 3), dtype int32. 25 | normals (torch.Tensor or str): 26 | GPU tensor of vertex normals. Tensor of shape (V, 3), dtype float32. 27 | Alternatively, you can specify 'flat' or 'smooth' here. 28 | Automatic normals with flat or smooth faces will be computed. 29 | Defaults to 'flat'. 30 | M (torch.Tensor): 31 | Model matrix, or pose transform of the object. 32 | Tensor of shape (4, 4), dtype float32. 33 | Defaults to identity. 34 | colors (torch.Tensor): 35 | Linear-space RGBA vertex colors, in range [0, 1]. 36 | Tensor of shape (V, 4), dtype float32. 37 | Defaults to white (all ones). 38 | uv (torch.Tensor): 39 | UV coordinates of vertices. Tensor of shape (V, 2), dtype float32. 40 | Defaults to zeros. 41 | tangents (torch.Tensor): 42 | Tangent spaces of vertices. Tensor of shape (V, 4), dtype float32. 43 | The first 3 dimensions are the tangent vector. 44 | The last dimension is sign of the bitangent. 45 | Defaults to zeros. 46 | custom_attrs (Dict[str, torch.Tensor]): 47 | Arbitrary attributes you want to bind to your vertices. 48 | Tensor of shape (V, \\*) for each attribute, dtype float32. 49 | metadata (Dict[str, Any]): 50 | Arbitrary meta data to tag the object. Not used for rendering. 51 | """ 52 | # material 53 | material: 'SurfaceMaterial' 54 | 55 | # geometry 56 | verts: torch.Tensor 57 | tris: torch.IntTensor 58 | normals: Union[Literal['flat', 'smooth'], torch.Tensor] = 'flat' 59 | 60 | # transform 61 | M: Optional[torch.Tensor] = None 62 | 63 | # attributes 64 | color: Optional[torch.Tensor] = None 65 | uv: Optional[torch.Tensor] = None 66 | tangents: Optional[torch.Tensor] = None 67 | custom_attrs: Optional[Dict[str, torch.Tensor]] = None 68 | 69 | metadata: Optional[Dict[str, Any]] = None 70 | 71 | def preprocess(self): 72 | if self.M is None: 73 | self.M = torch.eye(4, device=self.verts.device, dtype=self.verts.dtype) 74 | if self.color is None: 75 | self.color = ones_like_vec(self.verts, 4) 76 | if self.uv is None: 77 | self.uv = zeros_like_vec(self.verts, 2) 78 | if self.tangents is None: 79 | self.tangents = zeros_like_vec(self.verts, 4) 80 | if self.custom_attrs is None: 81 | self.custom_attrs = {} 82 | if self.metadata is None: 83 | self.metadata = {} 84 | if isinstance(self.normals, str): 85 | if self.normals == 'smooth': 86 | self.normals = compute_vertex_normals(self.verts, self.tris) 87 | if self.normals == 'flat': 88 | f = self.tris 89 | fn = compute_face_normals(self.verts, self.tris, normalize=True) 90 | self.verts, self.tris, self.normals = make_face_soup(self.verts, self.tris, fn) 91 | self.color = self.color[f].flatten(0, 1) 92 | self.uv = self.uv[f].flatten(0, 1) 93 | self.tangents = self.tangents[f].flatten(0, 1) 94 | new_customs = {} 95 | for k, attr in self.custom_attrs.items(): 96 | new_customs[k] = attr[f].flatten(0, 1) 97 | self.custom_attrs = new_customs 98 | return self 99 | -------------------------------------------------------------------------------- /diffrp/scene/scene.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import collections 3 | from typing import List 4 | from .lights import Light 5 | from .objects import MeshObject 6 | from typing_extensions import Self 7 | from ..utils.shader_ops import transform_point4x3, transform_vector3x3, float4, normalized 8 | 9 | 10 | class Scene: 11 | """ 12 | Represents a scene in DiffRP. 13 | ``add_*`` methods return the scene itself for chaining calls. 14 | 15 | Attributes: 16 | lights: List of light sources in the scene. 17 | objects: List of objects in the scene. Currently supporting MeshObject only. 18 | metadata: Dict of arbitrary loader-specific or user-defined metadata. 19 | """ 20 | def __init__(self) -> None: 21 | self.lights: List[Light] = [] 22 | self.objects: List[MeshObject] = [] 23 | self.metadata = {} 24 | 25 | def add_light(self, light: Light) -> Self: 26 | self.lights.append(light) 27 | return self 28 | 29 | def add_mesh_object(self, mesh_obj: MeshObject) -> Self: 30 | self.objects.append(mesh_obj.preprocess()) 31 | return self 32 | 33 | def static_batching(self) -> Self: 34 | """ 35 | Batches meshes with the same material into a single pass for faster rendering. 36 | Suitable if you have multiple meshes with the same material, and 37 | want to render multiple frames of the same scene. 38 | 39 | Changes are made in-place. 40 | You shall not modify the old mesh objects after calling this method. 41 | Local coordinates may change for individual meshes as they are combined. 42 | The world space and most renders will remain the same after the operation. 43 | 44 | Meshes with the same material are required to have the same set of custom attributes. 45 | 46 | Returns: 47 | Scene: The scene itself after in-place modification. 48 | """ 49 | material_list = collections.defaultdict(list) 50 | for mesh in self.objects: 51 | material_list[id(mesh.material)].append(mesh) 52 | 53 | def _combine(meshes: List[MeshObject]) -> MeshObject: 54 | if len(meshes) == 1: 55 | return meshes[0] 56 | assert len(meshes) > 1 57 | tris = [] 58 | offset = 0 59 | for x in meshes: 60 | tris.append(x.tris + offset) 61 | offset += len(x.verts) 62 | return MeshObject( 63 | x.material, 64 | torch.cat([transform_point4x3(x.verts, x.M) for x in meshes]), 65 | torch.cat(tris), 66 | torch.cat([normalized(transform_vector3x3(x.normals, x.M)) for x in meshes]), 67 | torch.eye(4, dtype=torch.float32, device=x.verts.device), 68 | torch.cat([x.color for x in meshes]), 69 | torch.cat([x.uv for x in meshes]), 70 | torch.cat([float4(normalized(transform_vector3x3(x.tangents.xyz, x.M)), x.tangents.w) for x in meshes]), 71 | {k: torch.cat([x.custom_attrs[k] for x in meshes]) for k in x.custom_attrs} 72 | ) 73 | 74 | self.objects = [_combine(x) for x in material_list.values()] 75 | return self 76 | -------------------------------------------------------------------------------- /diffrp/scripts/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/diffrp/scripts/__init__.py -------------------------------------------------------------------------------- /diffrp/scripts/main.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from .quick_render import build_argparse 3 | 4 | 5 | def main(): 6 | argp = argparse.ArgumentParser(prog='diffrp', description='diffrp CLI entrypoint.') 7 | subparsers = argp.add_subparsers(help='available diffrp scripts.', required=True) 8 | build_argparse(subparsers.add_parser("quick-render", help="Quickly render previews with the deferred rasterization PBR pipeline for 3D models.")) 9 | args = argp.parse_args() 10 | args.entrypoint(args) 11 | 12 | 13 | if __name__ == '__main__': 14 | main() 15 | -------------------------------------------------------------------------------- /diffrp/scripts/quick_render.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import diffrp 4 | import psutil 5 | import trimesh 6 | import rich.progress 7 | from diffrp.loaders.gltf_loader import _load_glb 8 | from diffrp.utils import * 9 | from diffrp.resources.hdris import newport_loft 10 | from multiprocessing.pool import ThreadPool 11 | 12 | 13 | @torch.no_grad() 14 | def normalize_to_2_unit_cube(scene: diffrp.Scene): 15 | # Get the bounding box of the mesh 16 | bmin = diffrp.gpu_f32([1e30] * 3) 17 | bmax = diffrp.gpu_f32([-1e30] * 3) 18 | world_v = [diffrp.transform_point4x3(prim.verts, prim.M) for prim in scene.objects] 19 | for verts, prim in zip(world_v, scene.objects): 20 | bmin = torch.minimum(bmin, verts.min(0)[0]) 21 | bmax = torch.maximum(bmax, verts.max(0)[0]) 22 | min_bounds = bmin 23 | max_bounds = bmax 24 | 25 | # Compute the center of the mesh 26 | center = ((min_bounds + max_bounds) / 2.0).cpu().numpy() 27 | 28 | # Compute the maximum dimension of the bounding box 29 | scale = torch.max(max_bounds - min_bounds).item() 30 | 31 | T = trimesh.transformations.translation_matrix(-center) 32 | S = trimesh.transformations.scale_matrix(2.0 / scale) 33 | M = diffrp.gpu_f32(S @ T) 34 | for prim in scene.objects: 35 | prim.M = M @ prim.M 36 | return scene 37 | 38 | 39 | def load(p: str): 40 | if p.endswith(".glb"): 41 | with open(p, "rb") as fi: 42 | return os.path.basename(p), trimesh.load(_load_glb(fi), force='scene', process=False) 43 | else: 44 | return os.path.basename(p), trimesh.load(p, force='scene', process=False) 45 | 46 | 47 | def build_argparse(argp: argparse.ArgumentParser): 48 | argp.set_defaults(entrypoint=main) 49 | argp.add_argument("-i", "--glb-path", required=True) 50 | argp.add_argument("-o", "--output-path", default='rendered') 51 | argp.add_argument("-t", "--io-threads", type=int, default=psutil.cpu_count(False)) 52 | argp.add_argument("-r", "--resolution", type=int, default=512) 53 | argp.add_argument("-s", "--ssaa", type=int, default=2) 54 | argp.add_argument("-f", "--fov", type=int, default=30) 55 | argp.add_argument("-a", "--azim-list", type=str, default='0,45,90,135,180,225,270,315') 56 | argp.add_argument("-e", "--elev-list", type=str, default='0,30,-15') 57 | 58 | 59 | @torch.no_grad() 60 | def main(args): 61 | todo = [] 62 | 63 | if os.path.isfile(args.glb_path): 64 | todo.append(args.glb_path) 65 | else: 66 | for r, ds, fs in os.walk(args.glb_path): 67 | for f in fs: 68 | if f.endswith('.glb'): 69 | todo.append(os.path.join(r, f)) 70 | write_jobs = [] 71 | light_cache = None 72 | os.makedirs(args.output_path, exist_ok=True) 73 | image_per_obj = len(args.elev_list.split(",")) * len(args.azim_list.split(",")) 74 | with ThreadPool(args.io_threads) as pool: 75 | with rich.progress.Progress() as prog: 76 | obj_task = prog.add_task("Rendering...", total=len(todo)) 77 | sub_task = prog.add_task("", total=image_per_obj) 78 | for name, scene in pool.imap_unordered(load, todo): 79 | prog.update(obj_task, advance=1) 80 | prog.update(sub_task, total=image_per_obj, completed=0, description=name) 81 | scene = diffrp.from_trimesh_scene(scene, compute_tangents=True).static_batching() 82 | scene = normalize_to_2_unit_cube(scene) 83 | cam_id = 0 84 | scene.add_light(diffrp.ImageEnvironmentLight(1.0, gpu_f32([1.0] * 3), newport_loft().cuda(), False)) 85 | for elev in map(float, args.elev_list.split(",")): 86 | for azim in map(float, args.azim_list.split(",")): 87 | R = args.resolution 88 | A = args.ssaa 89 | F = args.fov 90 | radius = 1.5 / math.sin(math.radians(F / 2)) 91 | camera = diffrp.PerspectiveCamera.from_orbit(R * A, R * A, radius, azim, elev, [0, 0, 0], F, far=20) 92 | cam_id += 1 93 | rp = diffrp.SurfaceDeferredRenderSession(scene, camera) 94 | if light_cache is None: 95 | light_cache = rp.prepare_ibl() 96 | else: 97 | rp.set_prepare_ibl(light_cache) 98 | pbr_premult = rp.compose_layers( 99 | rp.pbr_layered() + [torch.zeros([R * A, R * A, 3], device='cuda')], 100 | rp.alpha_layered() + [torch.zeros([R * A, R * A, 1], device='cuda')] 101 | ) 102 | pbr = ssaa_downscale(pbr_premult, 2) 103 | pbr = float4(agx_base_contrast(pbr.rgb / torch.clamp_min(pbr.a, 0.0001)), pbr.a) 104 | prog.update(sub_task, advance=1) 105 | write_jobs.append(pool.apply_async( 106 | to_pil(pbr).save, 107 | (os.path.join(args.output_path, f"{cam_id:02}_{name}.png"),), 108 | dict(compress_level=1) 109 | )) 110 | for job in rich.progress.track(write_jobs, "Writing..."): 111 | job.get() 112 | 113 | 114 | if __name__ == '__main__': 115 | main() 116 | -------------------------------------------------------------------------------- /diffrp/utils/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Includes various kind of utilities. 3 | """ 4 | 5 | from .colors import * 6 | from .composite import * 7 | from .coordinates import * 8 | from .exchange import * 9 | from .geometry import * 10 | from .light_transport import * 11 | from .shader_ops import * 12 | from .tone_mapping import * 13 | -------------------------------------------------------------------------------- /diffrp/utils/cache.py: -------------------------------------------------------------------------------- 1 | import functools 2 | from types import FunctionType 3 | from typing import Hashable, TypeVar 4 | 5 | 6 | class ICached: 7 | _cache: dict 8 | 9 | 10 | FuncType = TypeVar('FuncType') 11 | 12 | 13 | def cached_func(self: ICached, key: Hashable, func: FunctionType): 14 | if not hasattr(self, '_cache'): 15 | self._cache = {} 16 | if key not in self._cache: 17 | self._cache[key] = func() 18 | return self._cache[key] 19 | 20 | 21 | def cached(func: FuncType) -> FuncType: 22 | 23 | @functools.wraps(func) 24 | def wrapped(self: ICached): 25 | return cached_func(self, func.__qualname__, lambda: func(self)) 26 | 27 | return wrapped 28 | 29 | 30 | def key_cached(func: FuncType) -> FuncType: 31 | 32 | @functools.wraps(func) 33 | def wrapped(self: ICached, arg: Hashable): 34 | return cached_func(self, arg, lambda: func(self, arg)) 35 | 36 | return wrapped 37 | 38 | 39 | def singleton_cached(func: FuncType) -> FuncType: 40 | 41 | @functools.wraps(func) 42 | def wrapped(): 43 | return cached_func(func, func.__qualname__, lambda: func()) 44 | 45 | return wrapped 46 | -------------------------------------------------------------------------------- /diffrp/utils/colors.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .shader_ops import fsa 3 | 4 | 5 | PU_A = 1.41283765e+03 6 | PU_B = 1.64593172e+00 7 | PU_C = 4.31384981e-01 8 | PU_D = -2.94139609e-03 9 | PU_E = 1.92653254e-01 10 | PU_F = 6.26026094e-03 11 | PU_G = 9.98620152e-01 12 | PU_Y0 = 1.57945760e-06 13 | PU_Y1 = 3.22087631e-02 14 | PU_X0 = 2.23151711e-03 15 | PU_X1 = 3.70974749e-01 16 | 17 | 18 | def linear_to_pu(y): 19 | return torch.where(y <= PU_Y0, 20 | PU_A * y, 21 | torch.where(y <= PU_Y1, 22 | PU_B * torch.pow(y, PU_C) + PU_D, 23 | PU_E * torch.log(y + PU_F) + PU_G)) 24 | 25 | def pu_to_linear(x): 26 | return torch.where(x <= PU_X0, 27 | x / PU_A, 28 | torch.where(x <= PU_X1, 29 | torch.pow((x - PU_D) / PU_B, 1./PU_C), 30 | torch.exp((x - PU_G) / PU_E) - PU_F)) 31 | 32 | 33 | def linear_to_srgb(rgb: torch.Tensor) -> torch.Tensor: 34 | """ 35 | Converts from linear space to sRGB space. 36 | It is recommended to have the inputs already in LDR (range 0 to 1). 37 | 38 | The output has the same shape as the input. 39 | """ 40 | if rgb.requires_grad: 41 | rgb = torch.clamp_min(rgb, 1e-5) 42 | return torch.where(rgb < 0.0031308, 12.92 * rgb, fsa(1.055, rgb ** (1 / 2.4), -0.055)) 43 | 44 | 45 | def srgb_to_linear(rgb: torch.Tensor) -> torch.Tensor: 46 | """ 47 | Converts from sRGB space to linear space. 48 | 49 | The output has the same shape as the input. 50 | """ 51 | if rgb.requires_grad: 52 | rgb = torch.clamp_min(rgb, 1e-5) 53 | return torch.where(rgb < 0.04045, rgb * (1 / 12.92), (1 / 1.055) * ((rgb + 0.055) ** 2.4)) 54 | 55 | 56 | def aces_simple(x): 57 | """ 58 | Simple ACES mapping of color values. 59 | """ 60 | a = 2.51 61 | b = 0.03 62 | c = 2.43 63 | d = 0.59 64 | e = 0.14 65 | return torch.clamp((x * (a * x + b)) / (x * (c * x + d) + e), 0, 1) 66 | 67 | 68 | def aces_fit(color: torch.Tensor): 69 | """ 70 | More accurate ACES mapping of color values. 71 | """ 72 | aces_i = color.new_tensor([ 73 | [0.59719, 0.35458, 0.04823], 74 | [0.07600, 0.90834, 0.01566], 75 | [0.02840, 0.13383, 0.83777] 76 | ]) 77 | aces_o = color.new_tensor([ 78 | [ 1.60475, -0.53108, -0.07367], 79 | [-0.10208, 1.10813, -0.00605], 80 | [-0.00327, -0.07276, 1.07602] 81 | ]) 82 | 83 | v = color.matmul(aces_i.T) 84 | 85 | a = v * (v + 0.0245786) - 0.000090537 86 | b = v * (0.983729 * v + 0.4329510) + 0.238081 87 | v = a / b 88 | 89 | v = v.matmul(aces_o.T) 90 | 91 | return torch.clamp(v, 0, 1) 92 | 93 | 94 | def linear_to_alexa_logc_ei1000(x): 95 | cut = 0.010591 96 | a = 5.555556 97 | b = 0.052272 98 | c = 0.247190 99 | d = 0.385537 100 | e = 5.367655 101 | f = 0.092809 102 | return torch.where(x > cut, fsa(c, torch.log10(fsa(a, x, b)), d), fsa(e, x, f)) 103 | 104 | 105 | def alexa_logc_ei1000_to_linear(logc): 106 | cut = 0.010591 107 | a = 5.555556 108 | b = 0.052272 109 | c = 0.247190 110 | d = 0.385537 111 | e = 5.367655 112 | f = 0.092809 113 | return torch.where(logc > e * cut + f, (10 ** ((logc - d) / c) - b) / a, (logc - f) / e) 114 | -------------------------------------------------------------------------------- /diffrp/utils/composite.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Union, List 3 | from .shader_ops import saturate, split_alpha, to_bchw, to_hwc, fma 4 | 5 | 6 | def background_alpha_compose(bg: Union[torch.Tensor, List[Union[float, int]], float, int], fgca: torch.Tensor): 7 | """ 8 | Compose a background color to a foreground via alpha blending. 9 | 10 | It not only applies to RGB color and RGBA foregrounds, but also general linear data. 11 | 12 | The background should be one element shorter than the foreground elements. 13 | 14 | Args: 15 | bg: Tensor or List[float] or number. 16 | A single number is equivalent to the RGB list filled by it. 17 | The length should be one less than foreground. 18 | fgca: Tensor with alpha, shape (..., C + 1). 19 | Returns: 20 | torch.Tensor: Composed result, shape (..., C). 21 | """ 22 | fgc, fga = split_alpha(fgca) 23 | if not isinstance(bg, (float, int)) and not isinstance(bg, torch.Tensor): 24 | bg = fgca.new_tensor(bg) 25 | return fgc * fga + bg * (1 - fga) 26 | 27 | 28 | def alpha_blend(bgc: torch.Tensor, bga: torch.Tensor, fgc: torch.Tensor, fga: torch.Tensor): 29 | return torch.lerp(bgc, fgc, fga), fma(bga, 1 - fga, fga) 30 | 31 | 32 | def additive(bgc: torch.Tensor, bga: torch.Tensor, fgc: torch.Tensor, fga: torch.Tensor): 33 | return bgc + fgc, saturate(bga + fga) 34 | 35 | 36 | def alpha_additive(bgc: torch.Tensor, bga: torch.Tensor, fgc: torch.Tensor, fga: torch.Tensor): 37 | return bgc + fgc * fga, saturate(bga + fga) 38 | 39 | 40 | def ssaa_downscale(rgb, factor: int): 41 | return to_hwc(torch.nn.functional.interpolate(to_bchw(rgb), scale_factor=1 / factor, mode='area')) 42 | -------------------------------------------------------------------------------- /diffrp/utils/coordinates.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | from .shader_ops import flipper_2d 4 | 5 | 6 | def near_plane_ndc_grid(H: int, W: int, dtype: torch.dtype, device: torch.device): 7 | y = torch.linspace(-1 + 1 / H, 1 - 1 / H, H, dtype=dtype, device=device).view(H, 1, 1).expand(H, W, 1) 8 | x = torch.linspace(-1 + 1 / W, 1 - 1 / W, W, dtype=dtype, device=device).view(1, W, 1).expand(H, W, 1) 9 | near = (-flipper_2d()).expand(H, W, 2) 10 | return torch.cat([x, y, near], dim=-1) 11 | 12 | 13 | def latlong_grid(H, W, dtype, device): 14 | phi = torch.linspace(math.pi / 2, -math.pi / 2, H, dtype=dtype, device=device).view(-1, 1, 1) 15 | theta = torch.arange(W, dtype=dtype, device=device).view(-1, 1) * (math.tau / (W - 1)) 16 | return phi, theta 17 | 18 | 19 | def angles_to_unit_direction(theta: torch.Tensor, phi: torch.Tensor): 20 | """ 21 | Convert angle tensors (theta, phi) into unit direction vectors (x, y, z). 22 | 23 | The method is broadcastable. 24 | 25 | Inputs angles are in radians, unlike ``from_orbit``. 26 | 27 | See also: 28 | :py:obj:`diffrp.rendering.camera.PerspectiveCamera.from_orbit`. 29 | """ 30 | return _angles_to_unit_direction_impl(theta, phi) 31 | 32 | 33 | @torch.jit.script 34 | def _angles_to_unit_direction_impl(theta: torch.Tensor, phi: torch.Tensor): 35 | z = torch.cos(theta) * torch.cos(phi) 36 | x = torch.sin(theta) * torch.cos(phi) 37 | y = torch.sin(phi) 38 | return x, y, z 39 | 40 | 41 | def unit_direction_to_angles(L: torch.Tensor): 42 | """ 43 | Convert a unit direction vector (...batch_dims, 3) into two angle tensors (theta, phi). 44 | Theta and phi refer to the azimuth and elevation, respectively. 45 | 46 | Theta and phi will have one fewer dimension, shape (...batch_dims). 47 | These angles are in radians, unlike ``from_orbit``. 48 | 49 | See also: 50 | :py:obj:`diffrp.rendering.camera.PerspectiveCamera.from_orbit`. 51 | """ 52 | z = L.z 53 | if L.requires_grad: 54 | z = torch.where(z == 0, 1e-8, z) 55 | theta = torch.atan2(L.x, z) 56 | phi = torch.asin(torch.clamp(L.y, -0.999999, 0.999999)) 57 | return theta, phi 58 | 59 | 60 | @torch.jit.script 61 | def _unit_direction_to_latlong_uv_impl(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor): 62 | u = (torch.atan2(x, z) * (0.5 / math.pi)) % 1 63 | v = (1 / math.pi) * torch.asin(torch.clamp(y, -0.999999, 0.999999)) + 0.5 64 | return u, v 65 | 66 | 67 | def unit_direction_to_latlong_uv(L: torch.Tensor): 68 | z = L.z 69 | if L.requires_grad: 70 | z = torch.where(z == 0, 1e-8, z) 71 | return _unit_direction_to_latlong_uv_impl(L.x, L.y, z) 72 | -------------------------------------------------------------------------------- /diffrp/utils/exchange.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch_redstone as rst 3 | from PIL import Image 4 | from .shader_ops import saturate 5 | 6 | 7 | @torch.no_grad() 8 | def to_pil(rgb_or_rgba: torch.Tensor): 9 | """ 10 | Convert a RGB or RGBA tensor to an PIL Image. 11 | The tensor is assumed to have range [0, 1]. 12 | Values beyond that are clamped. 13 | 14 | Returns: 15 | Image: The converted PIL Image. 16 | """ 17 | img = rst.torch_to_numpy((saturate(rgb_or_rgba.float()) * 255).byte().contiguous()) 18 | return Image.fromarray(img) 19 | 20 | 21 | def torch_load_no_warning(fn: str): 22 | if torch.__version__ >= '2.4': 23 | return torch.load(fn, weights_only=True) 24 | else: 25 | return torch.load(fn) 26 | -------------------------------------------------------------------------------- /diffrp/utils/geometry.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .shader_ops import normalized, cross 3 | 4 | 5 | def make_face_soup(verts, tris, face_normals): 6 | """ 7 | Make a face soup geometry from vertices, triangles and face normals. 8 | 9 | Args: 10 | verts (torch.Tensor): Tensor of shape (num_verts, 3) containing vertex coordinates. 11 | faces (torch.Tensor): Tensor of shape (num_faces, 3) containing triangle indices. 12 | face_normals (torch.Tensor): Tensor of shape (num_faces, 3) containing face normals. 13 | 14 | Returns: 15 | tuple of (verts, tris, vertex_normals): 16 | Constructed geometry. 17 | There are 3 * num_faces vertices and num_faces triangles in the result. 18 | """ 19 | return ( 20 | verts[tris].reshape(-1, 3), 21 | torch.arange(len(tris.reshape(-1)), device=tris.device, dtype=torch.int32).reshape(tris.shape), 22 | torch.stack([face_normals] * 3, axis=1).reshape(-1, 3) 23 | ) 24 | 25 | 26 | def compute_face_normals(verts: torch.Tensor, faces: torch.Tensor, normalize=False): 27 | """ 28 | Compute per-face normals from vertices and faces. 29 | Faces are assumed to be CCW winding. 30 | 31 | Args: 32 | verts (torch.Tensor): Tensor of shape (num_verts, 3) containing vertex coordinates. 33 | faces (torch.Tensor): Tensor of shape (num_faces, 3) containing triangle indices. 34 | normalize (bool): Whether the results should be normalized. 35 | 36 | Returns: 37 | torch.Tensor: Tensor of shape (num_faces, 3) containing per-face normals. 38 | """ 39 | face_normals = cross(verts[faces[:, 1]] - verts[faces[:, 0]], 40 | verts[faces[:, 2]] - verts[faces[:, 0]]) 41 | return normalized(face_normals) if normalize else face_normals 42 | 43 | 44 | def compute_vertex_normals(verts: torch.Tensor, faces: torch.Tensor): 45 | """ 46 | Compute per-vertex normals from vertices and faces. 47 | Faces are assumed to be CCW winding. 48 | 49 | Args: 50 | verts (torch.Tensor): Tensor of shape (num_verts, 3) containing vertex coordinates. 51 | faces (torch.Tensor): Tensor of shape (num_faces, 3) containing triangle indices. 52 | 53 | Returns: 54 | torch.Tensor: Tensor of shape (num_verts, 3) containing normalized per-vertex normals. 55 | """ 56 | face_normals = compute_face_normals(verts, faces) 57 | 58 | vertex_normals = torch.zeros_like(verts) 59 | vertex_normals = vertex_normals.index_add(0, faces[:, 0], face_normals) 60 | vertex_normals = vertex_normals.index_add(0, faces[:, 1], face_normals) 61 | vertex_normals = vertex_normals.index_add(0, faces[:, 2], face_normals) 62 | 63 | vertex_normals = normalized(vertex_normals) 64 | 65 | return vertex_normals 66 | 67 | 68 | def sign2d(p1: torch.Tensor, p2: torch.Tensor, p3: torch.Tensor) -> torch.Tensor: 69 | """ 70 | Modified from ``calibur.sign2d``. 71 | """ 72 | return (p1.x - p3.x) * (p2.y - p3.y) - (p2.x - p3.x) * (p1.y - p3.y) 73 | 74 | 75 | def point_in_tri2d(pt: torch.Tensor, v1: torch.Tensor, v2: torch.Tensor, v3: torch.Tensor) -> torch.Tensor: 76 | """ 77 | Decide whether points lie in 2D triangles. Tensors broadcast. 78 | Modified from ``calibur.point_in_tri2d``. 79 | 80 | :param pt: ``(..., 2)`` points to decide. 81 | :param v1: ``(..., 2)`` first vertex in triangles. 82 | :param v2: ``(..., 2)`` second vertex in triangles. 83 | :param v3: ``(..., 2)`` third vertex in triangles. 84 | :returns: ``(..., 1)`` boolean result of points lie in triangles. 85 | """ 86 | d1 = sign2d(pt, v1, v2) 87 | d2 = sign2d(pt, v2, v3) 88 | d3 = sign2d(pt, v3, v1) 89 | has_neg = (d1 < 0) | (d2 < 0) | (d3 < 0) 90 | has_pos = (d1 > 0) | (d2 > 0) | (d3 > 0) 91 | return ~(has_neg & has_pos) 92 | 93 | 94 | @torch.jit.script 95 | def _barycentric_impl(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, p: torch.Tensor): 96 | v0 = b - a 97 | v1 = c - a 98 | v2 = p - a 99 | d00 = (v0 * v0).sum(-1) 100 | d01 = (v0 * v1).sum(-1) 101 | d11 = (v1 * v1).sum(-1) 102 | d20 = (v2 * v0).sum(-1) 103 | d21 = (v2 * v1).sum(-1) 104 | denom = d00 * d11 - d01 * d01 105 | v = (d11 * d20 - d01 * d21) / denom 106 | w = (d00 * d21 - d01 * d20) / denom 107 | v = torch.clip(torch.nan_to_num(v), 0, 1) 108 | w = torch.clip(torch.nan_to_num(w), 0, 1) 109 | u = 1 - v - w 110 | return torch.stack([u, v, w], dim=-1) 111 | 112 | 113 | def _cancel_nan(x): 114 | return torch.nan_to_num(x) 115 | 116 | 117 | def _cancel_nans(*tensors: torch.Tensor): 118 | for tnsr in tensors: 119 | if tnsr.requires_grad: 120 | tnsr.register_hook(_cancel_nan) 121 | 122 | 123 | def barycentric(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, p: torch.Tensor): 124 | """ 125 | Computes barycentric coordinates from a point ``p`` and triangle vertices ``a``, ``b`` and ``c``. 126 | 127 | Args: 128 | Takes broadcastable inputs of (..., 3). 129 | 130 | Returns: 131 | torch.Tensor: Stacked ``(u, v, w)`` tensor of shape (..., 3) s.t. ``p = ua + vb + wc``. 132 | """ 133 | _cancel_nans(a, b, c, p) 134 | return _barycentric_impl(a, b, c, p) 135 | -------------------------------------------------------------------------------- /diffrp/utils/light_transport.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn.functional as F 4 | from typing import Optional, Union 5 | from .cache import singleton_cached 6 | from .coordinates import angles_to_unit_direction, unit_direction_to_latlong_uv, latlong_grid 7 | from .shader_ops import normalized, float2, float3, cross, dot, sample2d, to_bchw, to_hwc, saturate, sample3d, fma 8 | 9 | 10 | @torch.jit.script 11 | def radical_inverse_van_der_corput(bits: torch.Tensor): 12 | bits = (bits << 16) | (bits >> 16) 13 | bits = ((bits & 0x55555555) << 1) | ((bits & 0xAAAAAAAA) >> 1) 14 | bits = ((bits & 0x33333333) << 2) | ((bits & 0xCCCCCCCC) >> 2) 15 | bits = ((bits & 0x0F0F0F0F) << 4) | ((bits & 0xF0F0F0F0) >> 4) 16 | bits = ((bits & 0x00FF00FF) << 8) | ((bits & 0xFF00FF00) >> 8) 17 | return bits.float() * 2.3283064365386963e-10 18 | 19 | 20 | def hammersley(n: int, deterministic: bool = False, device: Optional[torch.device] = None): 21 | return _hammersley_impl(n, deterministic, device) 22 | 23 | 24 | @torch.jit.script 25 | def _hammersley_impl(n: int, deterministic: bool, device: Optional[torch.device]): 26 | i = torch.arange(n, dtype=torch.int64, device=device) 27 | x = i.float() * (1 / n) 28 | y = radical_inverse_van_der_corput(i) 29 | if not deterministic: 30 | x = (x + torch.rand(1, dtype=x.dtype, device=x.device)) % 1.0 31 | y = (y + torch.rand(1, dtype=x.dtype, device=x.device)) % 1.0 32 | return x, y 33 | 34 | 35 | @torch.jit.script 36 | def combine_fixed_tangent_space(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor, n: torch.Tensor): 37 | up = torch.eye(3, dtype=n.dtype, device=n.device)[torch.where( 38 | n[..., 1] < 0.999, 1, 0 39 | )] 40 | # cannot use .y and .new_tensor shorthands for jit scripted function 41 | right = normalized(cross(up, n)) 42 | up = cross(n, right) 43 | return x * right + z * up + y * n 44 | 45 | 46 | def normal_distribution_function_ggx(n_dot_h: torch.Tensor, roughness: float): 47 | a = roughness * roughness 48 | f = n_dot_h * n_dot_h * (a * a - 1) + 1 49 | return (a * a / math.pi) / (f * f) 50 | 51 | 52 | def importance_sample_ggx(x: torch.Tensor, y: torch.Tensor, n: torch.Tensor, roughness: Union[float, torch.Tensor]): 53 | """ 54 | GGX Importance Sampling. Samples a batch of rays if roughness is a single float value, 55 | or samples a single ray for each (broadcastable) element if roughness is a tensor. 56 | 57 | Args: 58 | x (torch.Tensor): sample sequence element 1, shape (n) in [0, 1) if roughness is float, otherwise shape (..., 1) when roughness is tensor. 59 | y (torch.Tensor): sample sequence element 2, shape (n) in [0, 1) if roughness is float, otherwise shape (..., 1) when roughness is tensor. 60 | n (torch.Tensor): (batched) normal vectors, shape (..., 3). 61 | roughness (float | torch.Tensor): the roughness level, single value or tensor of shape (..., 1) in (0, 1]. 62 | 63 | Returns: 64 | torch.Tensor: sampled ray directions, shape (n, ..., 3) if roughness is float, otherwise shape (..., 3). 65 | """ 66 | if isinstance(roughness, float): 67 | return _importance_sample_ggx_impl_singler(x, y, n, roughness) 68 | else: 69 | return _importance_sample_ggx_impl_batchr(x, y, n, roughness) 70 | 71 | 72 | @torch.jit.script 73 | def _importance_sample_ggx_impl_batchr(x: torch.Tensor, y: torch.Tensor, n: torch.Tensor, roughness: torch.Tensor): 74 | a = roughness * roughness 75 | phi = math.tau * x 76 | cos_theta = torch.sqrt((1 - y) / (1 + (a * a - 1) * y)) 77 | sin_theta = torch.sqrt(1 - cos_theta * cos_theta) 78 | hx = torch.cos(phi) * sin_theta 79 | hz = torch.sin(phi) * sin_theta 80 | hy = cos_theta 81 | # n: ..., 3 82 | # h*: ..., 1 83 | return combine_fixed_tangent_space(hx, hy, hz, n) 84 | 85 | 86 | @torch.jit.script 87 | def _importance_sample_ggx_impl_singler(x: torch.Tensor, y: torch.Tensor, n: torch.Tensor, roughness: float): 88 | a = roughness * roughness 89 | phi = math.tau * x 90 | cos_theta = torch.sqrt((1 - y) / (1 + (a * a - 1) * y)) 91 | sin_theta = torch.sqrt(1 - cos_theta * cos_theta) 92 | hx = torch.cos(phi) * sin_theta 93 | hz = torch.sin(phi) * sin_theta 94 | hy = cos_theta 95 | # n: ..., 3 96 | # h*: n 97 | h_broadcast = [-1] + [1] * n.ndim # n, ...1, 1 98 | hx, hy, hz = [m.reshape(h_broadcast) for m in [hx, hy, hz]] 99 | return combine_fixed_tangent_space(hx, hy, hz, n) 100 | 101 | 102 | def prefilter_env_map( 103 | env: torch.Tensor, 104 | base_resolution: int = 256, num_levels: int = 5, 105 | num_samples: int = 512, deterministic: bool = False 106 | ): 107 | H = base_resolution 108 | W = H * 2 109 | env = to_hwc(F.interpolate(to_bchw(env), [H, W], mode='area')) 110 | pre_mips = [env] 111 | while H > 8: 112 | H = H // 2 113 | W = W // 2 114 | pre_mips.append(to_hwc(F.interpolate(to_bchw(pre_mips[-1]), [H, W], mode='area'))) 115 | H = base_resolution 116 | W = H * 2 117 | pre_mips = torch.stack([to_hwc(F.interpolate(to_bchw(mip), [H, W], mode='bilinear')) for mip in pre_mips]) 118 | levels = [env] 119 | rx, ry = hammersley(num_samples, deterministic, env.device) # n 120 | for i in range(1, num_levels): 121 | H = H // 2 122 | W = W // 2 123 | roughness = i / (num_levels - 1) 124 | phi, theta = latlong_grid(H, W, env.dtype, env.device) 125 | x, y, z = angles_to_unit_direction(theta, phi) 126 | r = float3(x, y, z) # h, w, 3 127 | 128 | n = v = r 129 | h = importance_sample_ggx(rx, ry, n, roughness) # n, h, w, 3 130 | ggx_d = normal_distribution_function_ggx(dot(n, h), roughness) # n, h, w, 1 131 | 132 | L = normalized(2.0 * dot(v, h) * h - v) # broadcast, n, h, w, 3 133 | del h 134 | n_dot_L = torch.relu_(dot(n, L)) # n, h, w, 1 135 | 136 | u, v = unit_direction_to_latlong_uv(L) 137 | satexel_inv = (2.0 * base_resolution * base_resolution) / math.pi 138 | sasample = (satexel_inv / num_samples) / (torch.clamp_min(1 - L.y ** 2, math.cos(math.pi / 2 - math.pi / 2 / H) ** 2) * (ggx_d + 0.0004)) 139 | del ggx_d, L 140 | 141 | mip_level = saturate(0.5 / (len(pre_mips) - 1) * torch.log2(sasample)) 142 | del sasample 143 | 144 | uv = float3(u, v, mip_level) 145 | del u, v, mip_level 146 | values = sample3d(pre_mips, uv, "reflection") # n, h, w, c 147 | del uv 148 | weights = n_dot_L.sum(0) 149 | 150 | values = torch.einsum('nhwc,nhw->hwc', values, n_dot_L.squeeze(-1)) 151 | values = values / (weights + 1e-4) 152 | values = to_hwc(F.interpolate(to_bchw(values), [base_resolution, base_resolution * 2], mode='bilinear', align_corners=True)) 153 | levels.append(values) 154 | levels.reverse() 155 | return torch.stack(levels) 156 | 157 | 158 | def irradiance_integral_env_map( 159 | env: torch.Tensor, 160 | premip_resolution: int = 64, 161 | base_resolution: int = 16 162 | ): 163 | HP, WP = premip_resolution, premip_resolution * 2 164 | env = to_hwc(F.interpolate(to_bchw(env), [HP, WP], mode='area')) 165 | H, W = base_resolution, base_resolution * 2 166 | phi, theta = latlong_grid(H, W, env.dtype, env.device) 167 | x, y, z = angles_to_unit_direction(theta, phi) 168 | n = float3(x, y, z)[..., None, None, :] # h, w, 3 -> h, w, 1, 1, 3 169 | sample_phi, sample_theta = latlong_grid(HP * 2, WP * 2, env.dtype, env.device) 170 | sample_phi = sample_phi[:len(sample_phi) // 2] 171 | x, y, z = angles_to_unit_direction(sample_theta, sample_phi) # tangent space, hp, wp, 1 172 | uv = float2(*unit_direction_to_latlong_uv(combine_fixed_tangent_space(x, y, z, n))) # h, w, hp, wp, 2 173 | values = sample2d(env, uv) * torch.cos(sample_phi) * torch.sin(sample_phi) # h, w, hp, wp, 3 174 | # sample_phi: hp, 1, 1 175 | values = values.mean([-2, -3]) 176 | return values * math.pi 177 | 178 | 179 | def geometry_schlick_ggx(n_dot_v: torch.Tensor, roughness: torch.Tensor): 180 | a = roughness 181 | k = (a * a) / 2.0 182 | nom = n_dot_v 183 | denom = n_dot_v * (1.0 - k) + k 184 | return nom / denom 185 | 186 | 187 | def geometry_smith(n: torch.Tensor, v: torch.Tensor, L: torch.Tensor, roughness: torch.Tensor): 188 | n_dot_v = torch.relu_(dot(n, v)) 189 | n_dot_L = torch.relu_(dot(n, L)) 190 | ggx2 = geometry_schlick_ggx(n_dot_v, roughness) 191 | ggx1 = geometry_schlick_ggx(n_dot_L, roughness) 192 | return ggx1 * ggx2 193 | 194 | 195 | def fresnel_schlick_smoothness(cos_theta: torch.Tensor, f0: torch.Tensor, smoothness: torch.Tensor): 196 | return fma(torch.relu_(smoothness - f0), (1.0 - cos_theta) ** 5.0, f0) 197 | 198 | 199 | @singleton_cached 200 | def pre_integral_env_brdf(): 201 | device = 'cuda' 202 | s = 256 203 | n_dot_v = torch.linspace(0 + 0.5 / s, 1 - 0.5 / s, s, device=device, dtype=torch.float32)[..., None] 204 | 205 | v = float3(torch.sqrt(1.0 - n_dot_v * n_dot_v), 0.0, n_dot_v) 206 | n = n_dot_v.new_tensor([0, 0, 1]).expand_as(v) 207 | 208 | rx, ry = hammersley(1024, True, device) # n 209 | results = [] 210 | for i in range(256): 211 | roughness = (i + 0.5) / 256 212 | h = importance_sample_ggx(rx, ry, n, roughness) # n, s, 3 213 | L = normalized(2 * dot(v, h) * h - v) 214 | # n_dot_L = torch.relu(L.z) 215 | n_dot_h = torch.relu(h.z) 216 | v_dot_h = torch.relu(dot(v, h)) 217 | g = geometry_smith(n, v, L, roughness) 218 | g_vis = (g * v_dot_h) / (n_dot_h * n_dot_v) 219 | fc = (1 - v_dot_h) ** 5 220 | results.append((g_vis * float2(1 - fc, fc)).mean(0)) 221 | return torch.stack(results) 222 | -------------------------------------------------------------------------------- /diffrp/utils/raycaster.py: -------------------------------------------------------------------------------- 1 | """ 2 | Experimental. Implementation and interface subject to change. 3 | """ 4 | import abc 5 | import torch 6 | from typing import Tuple 7 | from typing_extensions import Literal 8 | from torch_redstone import supercat 9 | 10 | from .shader_ops import cross, dot, normalized 11 | 12 | 13 | class Raycaster(metaclass=abc.ABCMeta): 14 | def __init__(self, verts: torch.Tensor, tris: torch.IntTensor, config: dict = {}) -> None: 15 | self.config = config 16 | self.build(verts, tris, config) 17 | 18 | @abc.abstractmethod 19 | def build(self, verts: torch.Tensor, tris: torch.IntTensor, config: dict) -> None: 20 | raise NotImplementedError 21 | 22 | @abc.abstractmethod 23 | def query(self, rays_o: torch.Tensor, rays_d: torch.Tensor, far: float) -> Tuple[torch.Tensor, torch.Tensor]: 24 | raise NotImplementedError 25 | 26 | 27 | def _ray_tri_pretransform(triangles: torch.Tensor): 28 | # M, 3, 3 29 | v0, v1, v2 = triangles.unbind(-2) 30 | e1 = v1 - v0 31 | e2 = v2 - v0 32 | nor = normalized(cross(e1, e2)) 33 | b2g = torch.stack([ 34 | v1 - v0, v2 - v0, nor, v0 35 | ], dim=-1) # M, 3, 4 36 | 37 | b2g = supercat([b2g, torch.eye(4, device=b2g.device)[-1]], dim=-2) 38 | g2b: torch.Tensor = torch.linalg.inv_ex(b2g)[0] 39 | return g2b[..., :-1, :-1], g2b[..., :-1, -1] 40 | 41 | 42 | @torch.jit.script 43 | def _ray_tri_intersect_pretransformed(rays_o: torch.Tensor, rays_d: torch.Tensor, far: float, g2br: torch.Tensor, g2bt: torch.Tensor): 44 | # ..., 1/M, 3 45 | # M, 3, 3 or ..., M, 3, 3 46 | # M, 3 or ..., M, 3 47 | rays_o = torch.matmul(g2br, rays_o.unsqueeze(-1)).squeeze(-1) + g2bt 48 | rays_d = torch.matmul(g2br, rays_d.unsqueeze(-1)).squeeze(-1) 49 | ox, oy, oz = torch.unbind(rays_o, dim=-1) 50 | dx, dy, dz = torch.unbind(rays_d, dim=-1) 51 | t = -oz / dz 52 | b1 = ox + t * dx 53 | b2 = oy + t * dy 54 | hit_mask = (t > 0) & (b1 >= 0) & (b2 >= 0) & (b1 + b2 <= 1) 55 | return torch.where(hit_mask, t, far) 56 | 57 | 58 | @torch.jit.script 59 | def _ray_tri_intersect( 60 | rays_o: torch.Tensor, rays_d: torch.Tensor, far: float, tris: torch.Tensor, epsilon: float 61 | ): 62 | # rays and tri verts should be broadcastable, and triangles should have at most one more dim 63 | # ..., 1/M, 3 64 | v1, v2, v0 = torch.unbind(tris, -2) # M, 3 or ..., M, 3 65 | e1 = v1 - v0 66 | e2 = v2 - v0 67 | cr = cross(rays_d, e2) 68 | det = dot(e1, cr, keepdim=False) 69 | inv_det = torch.reciprocal(det) 70 | s = rays_o - v0 71 | u = inv_det * dot(s, cr, keepdim=False) 72 | scross1 = cross(s, e1) 73 | v = inv_det * dot(rays_d, scross1, keepdim=False) 74 | t = inv_det * dot(e2, scross1, keepdim=False) # ..., M 75 | hit_mask = ( 76 | (torch.abs(det) > epsilon) & (u >= 0) & (v >= 0) & (u + v <= 1) & (t > 0) 77 | ) 78 | t = torch.where(hit_mask, t, far) # ..., M 79 | return t 80 | 81 | 82 | class BruteForceRaycaster(Raycaster): 83 | 84 | @torch.no_grad() 85 | def build(self, verts: torch.Tensor, tris: torch.IntTensor, config: dict): 86 | self.triangles = verts[tris] 87 | 88 | @torch.no_grad() 89 | def query(self, rays_o: torch.Tensor, rays_d: torch.Tensor, far: float): 90 | t = _ray_tri_intersect( 91 | rays_o.unsqueeze(-2), rays_d.unsqueeze(-2), far, self.triangles.unsqueeze(-4), 92 | self.config.get("epsilon", 1e-6) 93 | ) 94 | sel_tri = torch.argmin(t, dim=-1, keepdim=True) # ..., M -> ..., 1 95 | sel_t = torch.gather(t, -1, sel_tri).squeeze(-1) # ... 96 | fill_idx = sel_tri.squeeze(-1) 97 | return sel_t, fill_idx 98 | 99 | 100 | @torch.jit.script 101 | def _expand_bits(v: torch.Tensor): 102 | # long/int tensor 103 | v = (v * 0x00010001) & 0xFF0000FF 104 | v = (v * 0x00000101) & 0x0F00F00F 105 | v = (v * 0x00000011) & 0xC30C30C3 106 | v = (v * 0x00000005) & 0x49249249 107 | return v 108 | 109 | 110 | @torch.jit.script 111 | def zorder_3d(xyz: torch.Tensor): 112 | xyz = xyz - xyz.min(0).values 113 | xyz = xyz / (xyz.max(0).values + 1e-8) 114 | xyz = (xyz * 1023).int() 115 | xyz = _expand_bits(xyz) 116 | x, y, z = torch.unbind(xyz, dim=-1) 117 | return x | (y << 1) | (z << 2) 118 | 119 | 120 | class NaivePBBVH(Raycaster): 121 | 122 | @torch.no_grad() 123 | def build(self, verts: torch.Tensor, tris: torch.IntTensor, config: dict): 124 | M = len(tris) 125 | n = (M - 1).bit_length() 126 | 127 | triangles = verts[tris] 128 | 129 | builder: Literal['morton', 'splitaxis'] = config.get('builder', 'splitaxis') 130 | 131 | with torch.no_grad(): 132 | if M != 2 ** n: 133 | triangles = torch.cat([triangles, triangles[:2 ** n - M]]) 134 | centroids = triangles.mean(-2) 135 | if builder == 'morton': 136 | rank = torch.argsort(zorder_3d(centroids)) 137 | elif builder == 'splitaxis': 138 | rank = torch.arange(len(triangles), device=verts.device)[None] # 1, M 139 | smin = triangles.min(-2).values[None] # 1, ^M^, 3 140 | smax = triangles.max(-2).values[None] 141 | centroids = centroids[None] 142 | bmins = [] 143 | bmaxs = [] 144 | # 2 ** i, m, 3 145 | while smin.shape[-2] > 1: 146 | sbmin = smin.min(-2).values # 2 ** i, 3 147 | sbmax = smax.max(-2).values # 2 ** i, 3 148 | axes = torch.argmax(sbmax - sbmin, dim=-1) # 2 ** i 149 | i2, m = rank.shape 150 | b = torch.arange(i2, device=centroids.device) 151 | arg = torch.argsort(centroids[b, :, axes], dim=-1) # 2 ** i, m 152 | b = b.unsqueeze(-1) 153 | smin = smin[b, arg].view(2 * i2, m // 2, 3) 154 | smax = smax[b, arg].view(2 * i2, m // 2, 3) 155 | centroids = centroids[b, arg].view(2 * i2, m // 2, 3) 156 | rank = rank[b, arg].view(2 * i2, m // 2) 157 | rank = rank.view(2 ** n) 158 | elif builder == 'none': 159 | rank = torch.arange(len(triangles), device=verts.device) 160 | else: 161 | assert False, 'Unsupported builder: %s' % builder 162 | 163 | skip_next = torch.arange(2, (1 << (n + 1)) + 1, device=verts.device) 164 | skip_next = skip_next // ((-skip_next) & skip_next) - 1 165 | 166 | scan_next = torch.arange(1, (1 << (n + 2)) - 2, 2, device=verts.device) 167 | scan_next = torch.where(scan_next >= (1 << (n + 1)) - 1, skip_next, scan_next) 168 | 169 | bmins = [triangles.min(-2).values[rank]] 170 | bmaxs = [triangles.max(-2).values[rank]] 171 | while len(bmins[-1]) > 1: 172 | bmins.append(torch.minimum(*bmins[-1].view(-1, 2, 3).unbind(-2))) 173 | bmaxs.append(torch.maximum(*bmaxs[-1].view(-1, 2, 3).unbind(-2))) 174 | bmins.reverse() 175 | bmaxs.reverse() 176 | 177 | bmins = torch.cat(bmins) 178 | bmaxs = torch.cat(bmaxs) 179 | 180 | self.M = M 181 | self.n = n 182 | self.rank = rank % M 183 | self.bmins = bmins 184 | self.bmaxs = bmaxs 185 | self.skip_next = skip_next 186 | self.scan_next = scan_next 187 | self.g2br, self.g2bt = _ray_tri_pretransform(triangles[rank]) 188 | 189 | @staticmethod 190 | @torch.jit.script 191 | def _ray_box_intersect_impl( 192 | bmins: torch.Tensor, bmaxs: torch.Tensor, traverser: torch.Tensor, 193 | rays_o_live: torch.Tensor, rays_d_live: torch.Tensor, t_live: torch.Tensor, 194 | tri_start: int 195 | ): 196 | bound_min = bmins[traverser] 197 | bound_max = bmaxs[traverser] 198 | t1 = (bound_min - rays_o_live) / rays_d_live # R, 3 199 | t2 = (bound_max - rays_o_live) / rays_d_live # R, 3 200 | t_min = torch.min(t1, t2) # .max(-1).values 201 | t_max = torch.max(t1, t2) # .min(-1).values 202 | t_min = torch.max(torch.max(t_min[..., 0], t_min[..., 1]), t_min[..., 2]) 203 | t_max = torch.min(torch.min(t_max[..., 0], t_max[..., 1]), t_max[..., 2]) 204 | intersect_mask = (t_min <= t_live) & (t_max > 0) & (t_min <= t_max) 205 | return intersect_mask, intersect_mask & (traverser >= tri_start), traverser - tri_start 206 | 207 | @staticmethod 208 | @torch.jit.script 209 | def _traverse_tri_step_impl( 210 | live_idx: torch.Tensor, tri_hits: torch.Tensor, traverser_delta: torch.Tensor, far: float, 211 | rays_o: torch.Tensor, rays_d: torch.Tensor, g2br: torch.Tensor, g2bt: torch.Tensor, t: torch.Tensor, 212 | i: torch.Tensor 213 | ): 214 | ray_idx = live_idx[tri_hits] 215 | tri_idx = traverser_delta[tri_hits] 216 | # test_t = BruteForceRaycaster._ray_tri_intersect( 217 | # rays_o[ray_idx], rays_d[ray_idx], far, 218 | # self.triangles[tri_idx], epsilon 219 | # ) 220 | test_t = _ray_tri_intersect_pretransformed( 221 | rays_o[tri_hits], rays_d[tri_hits], far, g2br[tri_idx], g2bt[tri_idx] 222 | ) 223 | t.scatter_reduce_(0, ray_idx, test_t, 'amin') 224 | i[ray_idx] = torch.where(test_t <= t[ray_idx], tri_idx, i[ray_idx]) 225 | 226 | @torch.no_grad() 227 | def query(self, rays_o: torch.Tensor, rays_d: torch.Tensor, far: float): 228 | epsilon = self.config.get("epsilon", 1e-6) 229 | traverser = rays_o.new_zeros(rays_o.shape[:-1], dtype=torch.int64) # R 230 | # ^N^, 3 231 | # R, 3 232 | live_idx = torch.arange(len(traverser), device=traverser.device) 233 | t = rays_o.new_full(rays_o.shape[:-1], far) # R 234 | i = torch.zeros([len(rays_o)], device=traverser.device, dtype=torch.int64) 235 | tri_start = (1 << self.n) - 1 236 | while True: 237 | live_ray = traverser >= 0 238 | t_live = t[live_idx] 239 | for _ in range((self.n + 1) // 2): 240 | intersect_mask, tri_candidate, traverser_delta = NaivePBBVH._ray_box_intersect_impl( 241 | self.bmins, self.bmaxs, traverser, rays_o, rays_d, t_live, tri_start 242 | ) 243 | 244 | tri_hits = tri_candidate.nonzero().squeeze(-1) 245 | if len(tri_hits) > 0: 246 | NaivePBBVH._traverse_tri_step_impl( 247 | live_idx, tri_hits, traverser_delta, far, rays_o, rays_d, self.g2br, self.g2bt, t, i 248 | ) 249 | del tri_hits 250 | traverser = torch.where(intersect_mask, self.scan_next[traverser], self.skip_next[traverser]) 251 | live_ray &= traverser != 0 252 | live_ray = live_ray.nonzero(as_tuple=True) 253 | if len(live_ray[0]) == 0: 254 | break 255 | traverser = traverser[live_ray] 256 | live_idx = live_idx[live_ray] 257 | rays_o = rays_o[live_ray] 258 | rays_d = rays_d[live_ray] 259 | del live_ray 260 | return t, self.rank[i] 261 | 262 | 263 | class TorchOptiX(Raycaster): 264 | @torch.no_grad() 265 | def build(self, verts: torch.Tensor, tris: torch.IntTensor, config: dict) -> None: 266 | self.handle = None 267 | import torchoptix 268 | self.optix = torchoptix 269 | optix_log_level = config.get('optix_log_level', 3) 270 | self.optix.set_log_level(optix_log_level) 271 | self.verts = verts.contiguous() 272 | self.tris = tris.contiguous() 273 | self.handle = self.optix.build( 274 | self.verts.data_ptr(), self.tris.data_ptr(), 275 | len(verts), len(tris) 276 | ) 277 | 278 | @torch.no_grad() 279 | def query(self, rays_o: torch.Tensor, rays_d: torch.Tensor, far: float) -> Tuple[torch.Tensor]: 280 | out_t = rays_o.new_empty([len(rays_o)]) 281 | out_i = rays_o.new_empty([len(rays_o)], dtype=torch.int32) 282 | rays_o = rays_o.contiguous() 283 | rays_d = rays_d.contiguous() 284 | self.optix.trace_rays( 285 | self.handle, 286 | rays_o.data_ptr(), 287 | rays_d.data_ptr(), 288 | out_t.data_ptr(), out_i.data_ptr(), 289 | far, len(rays_o) 290 | ) 291 | return out_t, out_i 292 | 293 | def __del__(self): 294 | if self.handle is not None and self.optix is not None and self.optix.release is not None: 295 | self.optix.release(self.handle) 296 | self.handle = None 297 | -------------------------------------------------------------------------------- /diffrp/utils/tone_mapping.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .cache import key_cached 3 | from .shader_ops import sample3d 4 | from ..resources import get_resource_path 5 | from .colors import linear_to_alexa_logc_ei1000, linear_to_srgb 6 | from .exchange import torch_load_no_warning 7 | 8 | 9 | class AgxLutLoader: 10 | """ 11 | :meta private: 12 | """ 13 | @key_cached 14 | def load(self, variant: str) -> torch.Tensor: 15 | return torch.fliplr(torch_load_no_warning(get_resource_path("luts/agx-%s.pt" % variant)).cuda()).contiguous() 16 | 17 | 18 | agx_lut_loader = AgxLutLoader() 19 | 20 | 21 | def agx_base_contrast(rgb: torch.Tensor): 22 | """ 23 | Performs tone-mapping from HDR linear RGB space into LDR sRGB space 24 | with the state-of-the-art AgX tone-mapper, the base contrast variant. 25 | 26 | Args: 27 | rgb (torch.Tensor): Tensor of shape (..., 3) containing linear RGB values. 28 | Returns: 29 | torch.Tensor: Tensor of the same shape as the input. 30 | """ 31 | lut = agx_lut_loader.load("base-contrast") 32 | lut = lut.to(rgb) # z y x 3 33 | logc = linear_to_alexa_logc_ei1000(rgb) # h w 3 34 | # sampled: 1 c 1 h w 35 | return linear_to_srgb(sample3d(lut, logc)) 36 | -------------------------------------------------------------------------------- /diffrp/version.py: -------------------------------------------------------------------------------- 1 | __version__ = '0.2.7' 2 | -------------------------------------------------------------------------------- /diffrp/viewer/__init__.py: -------------------------------------------------------------------------------- 1 | from .mpl import Viewer 2 | -------------------------------------------------------------------------------- /diffrp/viewer/mpl.py: -------------------------------------------------------------------------------- 1 | """ 2 | Experimental. Implementation and interface subject to change. 3 | """ 4 | import time 5 | import numpy 6 | import torch 7 | import torch_redstone as rst 8 | from matplotlib import animation 9 | import matplotlib.pyplot as plotlib 10 | from ..rendering.camera import PerspectiveCamera 11 | from ..utils import background_alpha_compose, agx_base_contrast 12 | from ..rendering.surface_deferred import SurfaceDeferredRenderSession, SurfaceDeferredRenderSessionOptions 13 | 14 | 15 | class Viewer: 16 | """ 17 | View a scene by constructing the viewer, with the initial camera placed with the specified pose. 18 | Requires ``matplotlib`` to be installed. 19 | 20 | If you need semi-transparency you can tune up max_layers to the overdraw you want to support, 21 | but the performance would be lower. 22 | 23 | Use WASDQE to navigate around the scene, and rotate the view by holding the right mouse button. 24 | """ 25 | def __init__(self, scene, init_azim: int = 0, init_elev: int = 0, init_origin: list = [0, 0, 0], max_layers: int = 1): 26 | for key in plotlib.rcParams: 27 | if key.startswith('keymap.'): 28 | plotlib.rcParams[key] = [] 29 | self.scene = scene 30 | 31 | self.azim = init_azim 32 | self.elev = init_elev 33 | self.origin = numpy.array(init_origin.copy(), dtype=numpy.float32) 34 | self.key_state = set() 35 | self.sensitivity = 1.0 36 | self.max_layers = max_layers 37 | self.rebuild_camera() 38 | torch.set_grad_enabled(False) 39 | 40 | self.r_drag_pos = None 41 | self.ibl_cache = None 42 | self.last_update = None 43 | 44 | fig = plotlib.figure() 45 | ax: plotlib.Axes = plotlib.axes(navigate=False) 46 | ax.axis('off') 47 | self.widget = ax.imshow(rst.torch_to_numpy(self.render())) 48 | fig.subplots_adjust(left=0, right=1, bottom=0, top=1) 49 | fig.canvas.mpl_connect('button_press_event', self.on_press) 50 | fig.canvas.mpl_connect('button_release_event', self.on_release) 51 | fig.canvas.mpl_connect('motion_notify_event', self.on_motion) 52 | fig.canvas.mpl_connect('key_press_event', self.on_key_down) 53 | fig.canvas.mpl_connect('key_release_event', self.on_key_up) 54 | self.anim = animation.FuncAnimation(fig, self.refresh, interval=10, cache_frame_data=False, blit=True) 55 | plotlib.show() 56 | 57 | def update(self): 58 | t = time.perf_counter() 59 | if self.last_update is None: 60 | self.last_update = t 61 | delta = t - self.last_update 62 | velocity = numpy.zeros(3, dtype=numpy.float32) 63 | if 'a' in self.key_state: 64 | velocity[0] -= 1 65 | if 'd' in self.key_state: 66 | velocity[0] += 1 67 | if 'w' in self.key_state: 68 | velocity[2] -= 1 69 | if 's' in self.key_state: 70 | velocity[2] += 1 71 | if 'q' in self.key_state: 72 | velocity[1] -= 1 73 | if 'e' in self.key_state: 74 | velocity[1] += 1 75 | self.origin += velocity * delta * 5 76 | self.rebuild_camera() 77 | self.last_update = t 78 | 79 | def render(self): 80 | self.update() 81 | rp = SurfaceDeferredRenderSession( 82 | self.scene, self.camera, opaque_only=False, 83 | options=SurfaceDeferredRenderSessionOptions(max_layers=3) 84 | ) 85 | if self.ibl_cache is None: 86 | self.ibl_cache = rp.prepare_ibl() 87 | else: 88 | rp.set_prepare_ibl(self.ibl_cache) 89 | rgb = background_alpha_compose(1, rp.albedo()) 90 | return agx_base_contrast(rgb) 91 | 92 | def refresh(self, frame): 93 | import time 94 | a = time.perf_counter() 95 | render = self.render() 96 | b = time.perf_counter() 97 | render = rst.torch_to_numpy(render) 98 | c = time.perf_counter() 99 | self.widget.set_data(render) 100 | d = time.perf_counter() 101 | print("Render %.1f ms, Read %.1f ms, Set %.1f ms" % (1000 * (b - a), 1000 * (c - b), 1000 * (d - c))) 102 | return (self.widget,) 103 | 104 | def rebuild_camera(self): 105 | self.camera = PerspectiveCamera.from_orbit(512, 512, 3.0, self.azim, self.elev, self.origin, near=0.01, far=50.0) 106 | 107 | def on_key_down(self, event): 108 | self.key_state.add(event.key) 109 | 110 | def on_key_up(self, event): 111 | try: 112 | self.key_state.remove(event.key) 113 | except KeyError: 114 | import traceback; traceback.print_exc() 115 | 116 | def on_press(self, event): 117 | if event.button == 3: 118 | self.r_drag_pos = event.xdata, event.ydata 119 | 120 | def on_motion(self, event): 121 | if self.r_drag_pos is not None: 122 | oldx, oldy = self.r_drag_pos 123 | self.r_drag_pos = event.xdata, event.ydata 124 | self.azim -= event.xdata - oldx 125 | self.elev += event.ydata - oldy 126 | self.elev = min(90, max(self.elev, -90)) 127 | 128 | def on_release(self, event): 129 | if event.button == 3: 130 | self.r_drag_pos = None 131 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = source 9 | BUILDDIR = build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=source 11 | set BUILDDIR=build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.http://sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- 1 | sphinx==7.1.2 2 | furo==2023.9.10 3 | myst-parser==2.0.0 4 | matplotlib 5 | git+https://github.com/eliphatfs/diffrp.git 6 | -------------------------------------------------------------------------------- /docs/source/.gitignore: -------------------------------------------------------------------------------- 1 | generated -------------------------------------------------------------------------------- /docs/source/assets/game-1024spp-denoised.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/game-1024spp-denoised.jpg -------------------------------------------------------------------------------- /docs/source/assets/game-8spp-denoised-alpha.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/game-8spp-denoised-alpha.jpg -------------------------------------------------------------------------------- /docs/source/assets/game-8spp-noisy.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/game-8spp-noisy.jpg -------------------------------------------------------------------------------- /docs/source/assets/neural-albedo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/neural-albedo.png -------------------------------------------------------------------------------- /docs/source/assets/procedural-albedo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/procedural-albedo.png -------------------------------------------------------------------------------- /docs/source/assets/procedural-normals-small.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/procedural-normals-small.png -------------------------------------------------------------------------------- /docs/source/assets/procedural-normals.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/procedural-normals.png -------------------------------------------------------------------------------- /docs/source/assets/spheres-albedo-srgb.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/spheres-albedo-srgb.png -------------------------------------------------------------------------------- /docs/source/assets/spheres-attrs.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/spheres-attrs.jpg -------------------------------------------------------------------------------- /docs/source/assets/spheres-nvdraa-4xssaa.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/spheres-nvdraa-4xssaa.jpg -------------------------------------------------------------------------------- /docs/source/assets/spheres-output.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/spheres-output.jpg -------------------------------------------------------------------------------- /docs/source/assets/spheres-pbr-linear-hdr.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/spheres-pbr-linear-hdr.jpg -------------------------------------------------------------------------------- /docs/source/assets/spheres.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eliphatfs/diffrp/07d15c208510f0b9d430e0431d8264c9f6f98ac8/docs/source/assets/spheres.zip -------------------------------------------------------------------------------- /docs/source/concepts.md: -------------------------------------------------------------------------------- 1 | # Common Concepts 2 | 3 | ## Scene 4 | 5 | A *scene* includes objects and lights. 6 | The data structure can be found in {py:class}`diffrp.scene.scene.Scene`. 7 | 8 | Currently, only `MeshObject`s are supported for objects, and only `ImageEnvironmentLight`s are supported for lights. 9 | 10 | **Mesh Object** 11 | 12 | A `MeshObject` is a conventional triangle mesh formed by vertices and triangles, associated with a *material* and optionally vertex attributes including vertex normals, tangents, colors, UVs and other custom attributes. 13 | 14 | ## Material 15 | 16 | A *material* defines the appearance of an object. In DiffRP, materials are designed to be compatible across different render pipelines, though there is only one pipeline implemented currently. As only mesh objects are considered, there is only the `SurfaceMaterial` base that defines appearances of surfaces now, or more theoretically, how each point on the surfaces interact with light. 17 | 18 | Materials can be thought as functions defined on a point on the surface. The point may be associated with a triangle on the mesh, a world coordinate, and many extensive attributes like vertex normals, UVs and so on. How a point on the surface interacts with light is local -- it should not be affected by other points on the surface, or lights in the scene. If you need interaction with the scene outside the `SurfaceInput` and `SurfaceUniform` objects you have in the material function, it would in most cases be better that you rethink about it and do it by extending the render pipelines. 19 | 20 | ## Color Spaces 21 | 22 | Physically correct light interactions should be computed in the linear space since light is linear -- light is made of photons, and photons behave the same as each other. Linear means linear in number of photons, which is linear to the energy of light. 23 | 24 | Thus, color data is assumed linear in DiffRP, including those loaded by {py:mod}`diffrp.loaders.gltf_loader`. 25 | 26 | However, this is not how our eyes work. The perception of human vision is in the sRGB space, which is a different space than linear. Most LDR images (e.g. JPEG/PNG) are thus stored in the sRGB space on the computers. The mapping is commonly wrongly referred to as *gamma correction* but it should be a slightly different mapping. You can find it in {py:mod}`diffrp.utils.colors`. 27 | 28 | You may need to convert input images (like environment images or texture maps in JPEG/PNG) from sRGB to Linear before doing PBR, and convert the output images from Linear to sRGB before you view them. EXR/HDR images are usually already in linear space. 29 | 30 | If you are not doing PBR but only working on processing various attributes or fitting baked colors without considering interations with lights, color spaces are not that much -- theoretically, the correct interpolation requires linear spaces but usually the error is also close to invisible in sRGB. You just make sure your are in the same space, or in any space you expect. 31 | 32 | ## Tone Mapping 33 | 34 | There is no cap on the energy of light in real world. Thus, you usually get HDR images with maximum levels larger than 1 during PBR. To output it into a regular PNG image, you not only need to convert the color space from linear to sRGB, but also deal with these beyond-the-limit values. 35 | 36 | A naive way is to clip the values larger than 1.0. While this can work to some degree, we usually observe unnatural saturated and overexposed regions with this. By the way, this is usually the 'None' tone-mapping in rendering engines if you happen to see these options in unity or blender. 37 | 38 | The correct way is to apply a dedicated tone mapper. DiffRP provides a differentiable implementation of the state-of-the-art AgX tone mapper in {py:mod}`diffrp.utils.tone_mapping`. This usually results in much more realistic images than naive clipping. 39 | 40 | ## Visibility Gradient 41 | 42 | If you move a triangle by a small distance, colors on its edge can be affected, thus this effect should contribute to the gradient of the render. However, the edge of a triangle is a discontinuous region in the function, leading to difficulties in automatic differentiation. 43 | 44 | Visibility gradient is thus referring to techniques to take account of this contribution, and is special to differentiable rendering. 45 | 46 | The surface deferred render pipeline in DiffRP implements the Edge Gradient and the anti-aliasing approximation from nvdiffrast as visibility gradient generators. 47 | 48 | The path tracing pipeline in DiffRP does not officially support any visibility gradient yet. 49 | -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import inspect 3 | import os 4 | import sys 5 | sys.path.append("../../") 6 | 7 | import diffrp 8 | 9 | 10 | def abspath(rel): 11 | """ 12 | Take paths relative to the current file and 13 | convert them to absolute paths. 14 | 15 | Parameters 16 | ------------ 17 | rel : str 18 | Relative path, IE '../stuff' 19 | 20 | Returns 21 | ------------- 22 | abspath : str 23 | Absolute path, IE '/home/user/stuff' 24 | """ 25 | 26 | # current working directory 27 | cwd = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) 28 | return os.path.abspath(os.path.join(cwd, rel)) 29 | 30 | 31 | extensions = [ 32 | "sphinx.ext.napoleon", # numpy-style docstring 33 | "myst_parser", 34 | ] # allows markdown 35 | myst_all_links_external = True 36 | 37 | myst_enable_extensions = [ 38 | "dollarmath" 39 | ] 40 | 41 | # Add any paths that contain templates here, relative to this directory. 42 | templates_path = ["_templates"] 43 | 44 | # The suffix(es) of source filenames. 45 | source_suffix = { 46 | ".rst": "restructuredtext", 47 | ".txt": "markdown", 48 | ".md": "markdown", 49 | } 50 | 51 | # The master toctree document. 52 | master_doc = "index" 53 | 54 | # General information about the project. 55 | project = "diffrp" 56 | copyright = "2023, eliphatfs" 57 | author = "eliphatfs" 58 | 59 | # The version info for the project you're documenting, acts as replacement for 60 | # |version| and |release|, also used in various other places throughout the 61 | # built documents. 62 | # The full version, including alpha/beta/rc tags. 63 | release = diffrp.__version__ 64 | 65 | # The language for content autogenerated by Sphinx. Refer to documentation 66 | # for a list of supported languages. 67 | # This is also used if you do content translation via gettext catalogs. 68 | # Usually you set "language" from the command line for these cases. 69 | language = "en" 70 | 71 | # List of patterns, relative to source directory, that match files and 72 | # directories to ignore when looking for source files. 73 | # This patterns also effect to html_static_path and html_extra_path 74 | exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] 75 | 76 | # The name of the Pygments (syntax highlighting) style to use. 77 | pygments_style = "sphinx" 78 | 79 | # If true, `todo` and `todoList` produce output, else they produce nothing. 80 | todo_include_todos = False 81 | 82 | # -- Options for HTML output -------------------------------------- 83 | 84 | # The theme to use for HTML and HTML Help pages 85 | html_theme = "furo" 86 | 87 | # options for furo 88 | html_theme_options = {} 89 | 90 | # Add any paths that contain custom static files (such as style sheets) here, 91 | # relative to this directory. They are copied after the builtin static files, 92 | # so a file named "default.css" will overwrite the builtin "default.css". 93 | # html_static_path = ["static"] 94 | # html_logo = "images/trimesh-logo.png" 95 | 96 | # custom css 97 | # html_css_files = ["custom.css"] 98 | 99 | html_context = { 100 | "display_github": True, 101 | "github_user": "eliphatfs", 102 | "github_repo": "diffrp", 103 | "github_version": "main", 104 | "conf_py_path": "/docs/source/", 105 | } 106 | 107 | # Output file base name for HTML help builder. 108 | htmlhelp_basename = "diffrpdoc" 109 | 110 | autodoc_default_options = { 111 | "autosummary": True, 112 | "special-members": "__init__", 113 | } 114 | -------------------------------------------------------------------------------- /docs/source/gltf.md: -------------------------------------------------------------------------------- 1 | # Rendering glTF (*.glb) files 2 | 3 | Short for *GL Transmission Format*, [glTF](https://github.com/KhronosGroup/glTF) has become a standardized format choice of 3D content. DiffRP implements a {py:class}`GLTFMaterial` that covers all features of the glTF 2.0 specification. 4 | 5 | In this section, we will render the [MetalRoughSpheresNoTextures](https://github.com/KhronosGroup/glTF-Sample-Assets/blob/main/Models/MetalRoughSpheresNoTextures/README.md) scene from glTF sample asset library to walk through basic rendering features of DiffRP. You can download the scene from the link above or as a zip file {download}`here`. 6 | 7 | We will use the surface deferred (rasterization) rendering pipeline for this tutorial. 8 | 9 | ## Steps 10 | 11 | ### 1. Loading the Scene 12 | 13 | DiffRP includes an easy method to load a GLB file into a scene. 14 | 15 | ```python 16 | import diffrp 17 | from diffrp.utils import * 18 | 19 | scene = diffrp.load_gltf_scene("spheres.glb") 20 | ``` 21 | 22 | If you want to support normal maps, you will need to compute tangents with `compute_tangents=True`. This would require `gcc` to be installed, 23 | which is usually already fulfilled in most Linux and Mac distributions. 24 | For Windows the Strawberry Perl [(https://strawberryperl.com/)](https://strawberryperl.com/) distribution of `gcc` would be recommended. 25 | 26 | ```python 27 | scene = diffrp.load_gltf_scene("spheres.glb", compute_tangents=True) 28 | ``` 29 | 30 | Optionally, if you plan to render many frames with this scene, you can compute static batching to batch sub-objects with the same material into one to make future renders more efficient. 31 | 32 | ```python 33 | scene = scene.static_batching() 34 | ``` 35 | 36 | ### 2. Creating a Camera 37 | 38 | We can create an orbital camera within DiffRP to look at the center of the scene. The scene is small so we place the camera quite close to the spheres. 39 | 40 | ```python 41 | camera = diffrp.PerspectiveCamera.from_orbit( 42 | h=1080, w=1920, # resolution 43 | radius=0.02, azim=0, elev=0, # orbit camera position, azim/elev in degrees 44 | origin=[0.003, 0.003, 0], # orbit camera focus point 45 | fov=30, near=0.005, far=10.0 # intrinsics 46 | ) 47 | ``` 48 | 49 | You can also inherit and implement the {py:class}`Camera` class to implement your own camera data provider. You need to implement the `V()` and `P()` methods to provide GPU tensors of the projection matrix and the view matrix **in OpenGL convention**. 50 | 51 | ### 3. Create the Render Session 52 | 53 | Now let's create the render session that is the core for rendering in DiffRP. 54 | 55 | Rendering sessions are designed to be short-lived in DiffRP. Each frame that has a different scene or camera setup should have its own new render session. 56 | 57 | ```python 58 | rp = diffrp.SurfaceDeferredRenderSession(scene, camera) 59 | ``` 60 | 61 | ### 4. Rendering Normals and BRDF Parameters 62 | 63 | Now that we have the render session, rendering is very straightforward: 64 | 65 | ```python 66 | nor = rp.false_color_camera_space_normal() 67 | mso = rp.false_color_mask_mso() 68 | ``` 69 | 70 | `mso` represents to **m**etallic, **s**moothness and ambient **o**cclusion attributes in a common PBR BRDF material. These are also defined in the glTF specification. 71 | 72 | The outputs are tensors of shape $[H, W, 4]$ that linearly represent the attributes, in RGBA order, and RGB represents XYZ or MSO in the case, respectively. Normals are mapped from $[-1, 1]$ into $[0, 1]$ so the results are valid LDR RGB images. Alpha is transparency. 73 | 74 | You can also obtain the raw normal vectors without mapping to RGB values: 75 | 76 | ```python 77 | rp.camera_space_normal() 78 | ``` 79 | 80 | As a syntactic sugar injected by DiffRP, you can obtain the `xyz` or `rgb` channels in a tensor by simply using the `.` syntax: 81 | 82 | ```python 83 | nor.xyz 84 | nor.rgb 85 | nor.a 86 | nor.rgb * nor.a 87 | ``` 88 | 89 | Single channel dims are kept for broadcastability. 90 | 91 | Note however the values are undefined in transparent regions. You may would like to composite with a background color if you need it: 92 | 93 | ```python 94 | background_alpha_compose([0.5, 0.5, 0.5], nor) 95 | background_alpha_compose(0.5, nor) 96 | ``` 97 | 98 | Both of these commands would place a gray background for the normal map, which refers to a zero normal vector. 99 | 100 | ### 5. Reading the Image 101 | 102 | If you need to save or export the renders, DiffRP provides a utility method `to_pil` to convert the results from GPU Tensors into PIL Images. Only LDR RGB/RGBA images are supported. 103 | 104 | ```python 105 | to_pil(nor).show() 106 | to_pil(mso).show() 107 | ``` 108 | 109 | 110 | ```{figure} assets/spheres-attrs.jpg 111 | :scale: 45% 112 | :alt: Spheres attributes rendered. 113 | 114 | *A light gray background is added for more clarity of the MSO attributes.* 115 | ``` 116 | 117 | ### 6. Rendering Albedo or Base Color 118 | 119 | Rendering the albedo or base color of the method is just as straightforward as other parameters. However, you need to take care of the color space. 120 | 121 | To be physically correct, all operations on True Colors, including albedo/base color, are in linear space. 122 | 123 | ```python 124 | albedo = rp.albedo() # albedo in linear space 125 | ``` 126 | 127 | If you need to output the albedo in sRGB space (the space your PNG files are in), you need to convert it. DiffRP can do it for you: 128 | 129 | ```python 130 | albedo_srgb = rp.albedo_srgb() 131 | to_pil(albedo_srgb).show() 132 | ``` 133 | 134 | You can also call the lower level color space utility yourself: 135 | 136 | ```python 137 | float4(linear_to_srgb(albedo.rgb), albedo.a) 138 | ``` 139 | 140 | More color space utilities can be found in {py:mod}`diffrp.utils.colors`. 141 | 142 | The result should look like: 143 | 144 | ```{figure} assets/spheres-albedo-srgb.png 145 | :scale: 45 % 146 | :alt: Spheres Albedo render. 147 | ``` 148 | 149 | ### 7. Other Attributes 150 | 151 | DiffRP also supports other attributes that may or may not be represented as an LDR RGB image. 152 | Here we list these useful methods including those we cover in other sections of this tutorial. 153 | Please see the relevant API reference in {py:class}`diffrp.rendering.surface_deferred.SurfaceDeferredRenderSession` for details explaining the attributes. 154 | 155 | The methods all return (H, W, C) full-screen renders. 156 | 157 | ```python 158 | rp.aov('my_aov', [0.]) 159 | rp.albedo() 160 | rp.albedo_srgb() 161 | rp.emission() 162 | rp.pbr() 163 | rp.false_color_camera_space_normal() 164 | rp.false_color_world_space_normal() 165 | rp.false_color_mask_mso() 166 | rp.false_color_nocs() 167 | rp.camera_space_normal() 168 | rp.world_space_normal() 169 | rp.local_position() 170 | rp.world_position() 171 | rp.depth() 172 | rp.distance() 173 | rp.view_dir() 174 | ``` 175 | 176 | ### 8. PBR Rendering 177 | 178 | Currently, DiffRP supports fully differentiable Image-based lighting in the deferred shading pipeline. The support for directional and point lights will be added soon. 179 | 180 | Let's add an DiffRP built-in HDRI as the ambient lighting into our scene: 181 | 182 | ```python 183 | import torch 184 | from diffrp.resources.hdris import newport_loft 185 | 186 | scene.add_light(diffrp.ImageEnvironmentLight(intensity=1.0, color=torch.ones(3, device='cuda'), image=newport_loft().cuda())) 187 | ``` 188 | 189 | A tint color can be specified in the `color` parameter as you like. Note that the tensors needs to be on GPUs. 190 | 191 | ```{note} 192 | You can use your own HDRI environment. DiffRP accepts images in Lat-Long format (usually the width is double the height). 193 | 194 | If you load it from a PNG/JPG file, you usually need to convert it from sRGB to Linear space. 195 | ``` 196 | 197 | Now recreate the render session as the scene has been changed: 198 | 199 | ```python 200 | rp = diffrp.SurfaceDeferredRenderSession(scene, camera) 201 | ``` 202 | 203 | Now we can issue the rendering call just as simple as before: 204 | 205 | ```python 206 | pbr = rp.pbr() 207 | to_pil(pbr).show() 208 | ``` 209 | 210 | ```{figure} assets/spheres-pbr-linear-hdr.jpg 211 | :scale: 25 % 212 | :alt: PBR raw HDR output in linear space. 213 | ``` 214 | 215 | ### 9. Post-processing: Tone-mapping and Anti-aliasing 216 | 217 | The previous image seems weirdly dark. It is due to we are in HDR linear space, but we are viewing it as a LDR sRGB PNG image, which is the wrong color space. 218 | 219 | The procedure of converting an HDR image into LDR is called tone-mapping. DiffRP efficiently implements the state-of-the-art [AgX](https://github.com/sobotka/AgX) tone-mapper in a differentiable manner. 220 | 221 | ```python 222 | pbr = diffrp.agx_base_contrast(pbr.rgb) 223 | ``` 224 | 225 | The result would be a valid RGB image in sRGB space. 226 | 227 | We are simply discarding the alpha channel as we have the environment map as the background. You can specify `render_skybox=False` in the `ImageEnvironmentLight` to render a transparent-film PBR image. You will then have to keep the alpha channel (see the example code for albedo) or compose with your background (see the example code for normal map) during tone-mapping. 228 | 229 | At this final stage, we may also apply anti-aliasing. 230 | 231 | DiffRP integrates the AA operation in `nvdiffrast` to provide visibility-based gradients on the geometry. 232 | 233 | ```python 234 | pbr_aa = rp.nvdr_antialias(pbr) 235 | to_pil(pbr_aa).show() 236 | ``` 237 | 238 | The order of anti-aliasing and tone-mapping leaves for your choice. 239 | 240 | If you aim for higher rendering quality, you can also use a simple SSAA technique. That is, specify a higher resolution when rendering, and downscale at the end. 241 | 242 | ```python 243 | pbr = ssaa_downscale(pbr, 2) 244 | ``` 245 | 246 | Note that anti-aliasing operations do not currently support transparency. 247 | 248 | 249 | ### 9+. Transparent Film 250 | 251 | To disable the background to obtain a useful alpha mask, you can use the `render_skybox` parameter when adding environment light: 252 | 253 | ```python 254 | scene.add_light(diffrp.ImageEnvironmentLight(intensity=1.0, color=torch.ones(3, device='cuda'), image=newport_loft().cuda(), render_skybox=False)) 255 | ``` 256 | 257 | During tone mapping, you need to keep the alpha channel unchanged. 258 | 259 | ```python 260 | pbr = rp.pbr() 261 | pbr = float4(agx_base_contrast(pbr.rgb), pbr.a) 262 | ``` 263 | 264 | After that, you may also compose with a white background if you like: 265 | 266 | ```python 267 | pbr = diffrp.background_alpha_compose(1, pbr) # input pbr 4 channels rgba, output 3 channels rgb 268 | ``` 269 | 270 | 271 | ### 10. Handling Semi-Transparency (Optional) 272 | 273 | Theoretically, alpha blending is not a physical behavior. It is an inaccurate approximation of transmission and refraction. 274 | Thus, it is usually counter-intuitive and tricky to handle semi-transparency in PBR rasterization pipelines. 275 | 276 | Firstly, you need to disable opaque geometry only mode 277 | (this is made the default as rendering semi-transparency is much more compute and memory-intensive, and is rarely used for differentiable rasterization): 278 | 279 | ```python 280 | rp = diffrp.SurfaceDeferredRenderSession(scene, camera, opaque_only=False) 281 | ``` 282 | 283 | This makes the render pipeline support semi-transparency, and it works if you have a environment background image. 284 | 285 | For transparent films, you may also go with the simple handling described in section 9+, but we may do a little bit better. 286 | A more delicate handling of semi-transparency in terms of alpha values will require deeper analysis into the pipeline. 287 | 288 | For PBR rendering, HDR radiance values are composed, instead of LDR colors. 289 | This is the tricky part as alpha-blending itself makes no physical sense for HDR radiance values. 290 | 291 | We can however assume that alpha linearly dims the radiance values (note that this assumption makes some sense but is not physically accurate), 292 | which makes the blending consistent with pre-multiplied alpha in traditional alpha-blending. 293 | Pre-multiplied means that the RGB values are multiplied with alpha when stored, which is already the final color contribution. 294 | This is in contrast to straight alpha which means RGB values are the raw colors, which is usually what happens in PNG files and other storage formats. 295 | By default, DiffRP also treats alpha as straight alpha. 296 | 297 | It is also easy to do pre-multiplied alpha blending with DiffRP instead. 298 | We just need to start with zeros in the color buffer when composing instead of starting with values in the last buffer. 299 | 300 | ```python 301 | pbr_premult = rp.compose_layers( 302 | rp.pbr_layered() + [torch.zeros([1080, 1920, 3], device='cuda')], 303 | rp.alpha_layered() + [torch.zeros([1080, 1920, 1], device='cuda')] 304 | ) 305 | ``` 306 | 307 | The RGB values now represent a more accurate radiance value according to the assumption (that alpha linearly dims radiance). 308 | However, the value is "pre-multiplied" and makes no sense if combined with a straight alpha, 309 | so this is not made the default behavior in the `SurfaceDeferredRenderSession`. 310 | To display the image in viewers with the expected colors, you need to convert premultiplied alpha to straight alpha after tone mapping. 311 | You may also want to saturate the colors (clamping into [0, 1]) to keep it valid in LDR space. 312 | 313 | ```python 314 | pbr = ssaa_downscale(pbr_premult, 2) 315 | pbr = saturate(float4(agx_base_contrast(pbr.rgb) / torch.clamp_min(pbr.a, 0.0001), pbr.a)) 316 | ``` 317 | 318 | However, the colors will be slightly off if the colors are actually clamped (larger than 1 before clamping). 319 | This is a limitation of traditional LDR images with straight alpha. 320 | 321 | You can also convert the premultiplied alpha to straight alpha before tone mapping. 322 | There is an approximation in the tone mapper this way as the input is no longer a physical radiance value for semi-transparent regions, 323 | but you won't introduce errors when clamping the outputs. 324 | 325 | ```python 326 | pbr = float4(agx_base_contrast(pbr.rgb / torch.clamp_min(pbr.a, 0.0001)), pbr.a) 327 | ``` 328 | 329 | Note that in both ways, you should handle anti-aliasing in premultiplied space. 330 | 331 | Also, to put it again, alpha blending is a non-physical approximation to the transmission and refraction of light. 332 | Thus, there is no "correct way" to handle alpha (other than extreme values of 0 and 1) together with a PBR pipeline. 333 | 334 | All of the above ways make different assumptions for semi-transparency, and all of them are correct if all alpha values are either 0 or 1, 335 | including the simple method described in Section 9+. 336 | 337 | 338 | ## Complete Example 339 | 340 | ![Example output](assets/spheres-output.jpg) 341 | 342 | ```python 343 | import torch 344 | from diffrp.resources import hdris 345 | import diffrp 346 | from diffrp.utils import * 347 | 348 | 349 | scene = diffrp.load_gltf_scene("spheres.glb") 350 | scene.add_light(diffrp.ImageEnvironmentLight( 351 | intensity=1.0, color=torch.ones(3, device='cuda'), 352 | image=hdris.newport_loft().cuda() 353 | )) 354 | 355 | camera = diffrp.PerspectiveCamera.from_orbit( 356 | h=1080, w=1920, # resolution 357 | radius=0.02, azim=0, elev=0, # orbit camera position, azim/elev in degrees 358 | origin=[0.003, 0.003, 0], # orbit camera focus point 359 | fov=30, near=0.005, far=10.0 # intrinsics 360 | ) 361 | 362 | rp = diffrp.SurfaceDeferredRenderSession(scene, camera) 363 | pbr = rp.pbr() 364 | 365 | pbr_aa = rp.nvdr_antialias(pbr.rgb) 366 | to_pil(agx_base_contrast(pbr_aa)).save("sphere-output.jpg") 367 | ``` 368 | -------------------------------------------------------------------------------- /docs/source/index.md: -------------------------------------------------------------------------------- 1 | # Welcome to DiffRP⚙️! 2 | 3 | **DiffRP** aims to provide an easy-to-use programming interface for rendering and differentiable rendering pipelines. 4 | 5 | ```{figure} assets/spheres-nvdraa-4xssaa.jpg 6 | :scale: 40 % 7 | :alt: Sample Render with DiffRP PBR Rasterization. 8 | 9 | The glTF [MetalRoughSpheresNoTextures](https://github.com/KhronosGroup/glTF-Sample-Assets/blob/main/Models/MetalRoughSpheresNoTextures/README.md) scene rendered with DiffRP PBR rasterization. 10 | ``` 11 | 12 | 13 | 14 | ```{figure} assets/game-1024spp-denoised.jpg 15 | :scale: 60 % 16 | :alt: Sample Render with DiffRP Path Tracing. 17 | 18 | The glTF [ABeautifulGame](https://github.com/KhronosGroup/glTF-Sample-Assets/tree/main/Models/ABeautifulGame) scene rendered with DiffRP Path Tracing. 19 | ``` 20 | 21 | ## Installation 22 | 23 | The package can be installed by: 24 | 25 | ```bash 26 | pip install diffrp 27 | ``` 28 | 29 | Or for latest development: 30 | 31 | ```bash 32 | pip install git+https://github.com/eliphatfs/diffrp 33 | ``` 34 | 35 | ```{note} 36 | DiffRP depends on PyTorch (`torch`). The default version `pip` resolves to may not come with the `cuda` version you want. It is recommended to install [PyTorch](https://pytorch.org/get-started/locally/#start-locally) before you install DiffRP so you can choose the version you like. 37 | ``` 38 | 39 | ```{note} 40 | DiffRP rendering is based on CUDA GPUs. You can develop without one, but a CUDA GPU is required to run the code. 41 | ``` 42 | 43 | ### Other Dependencies 44 | 45 | If you use rasterization in DiffRP, you need to have the CUDA development kit set up as we use the `nvdiffrast` backend. See also [https://nvlabs.github.io/nvdiffrast/#installation](https://nvlabs.github.io/nvdiffrast/#installation). 46 | 47 | If you want to use hardware ray tracing in DiffRP, you need to install `torchoptix` by `pip install torchoptix`. See also the hardware and driver requirements [https://github.com/eliphatfs/torchoptix?tab=readme-ov-file#requirements](https://github.com/eliphatfs/torchoptix?tab=readme-ov-file#requirements). 48 | 49 | If you plan on using plugins in DiffRP (currently when you compute tangents), `gcc` is required in path. This is already fulfilled in most Linux and Mac distributions. For Windows I recommend the Strawberry Perl [(https://strawberryperl.com/)](https://strawberryperl.com/) distribution of `gcc`. 50 | 51 | ## Get Started 52 | 53 | Rendering attributes of a mesh programmingly is incredibly simple with DiffRP compared to conventional render engines like Blender and Unity. 54 | 55 | In this section, we will get started to render a simple procedural mesh. 56 | 57 | First, let's import what we would need: 58 | 59 | ```python 60 | import diffrp 61 | import trimesh.creation 62 | from diffrp.utils import * 63 | ``` 64 | 65 | We now create a icosphere and render the camera space normal map, where RGB color values in $[0, 1]$ map linearly to $[-1, 1]$ in normal XYZ vectors. 66 | 67 | The first run may take some time due to compiling and importing the `nvdiffrast` backend. Following calls would be fast. 68 | 69 | ```python 70 | # create the mesh (cpu) 71 | mesh = trimesh.creation.icosphere(radius=0.8) 72 | # initialize the DiffRP scene 73 | scene = diffrp.Scene() 74 | # register the mesh, load vertices and faces arrays to GPU 75 | scene.add_mesh_object(diffrp.MeshObject(diffrp.DefaultMaterial(), gpu_f32(mesh.vertices), gpu_i32(mesh.faces))) 76 | # default camera at [0, 0, 3.2] looking backwards 77 | camera = diffrp.PerspectiveCamera() 78 | # create the SurfaceDeferredRenderSession, a deferred-rendering rasterization pipeline session 79 | rp = diffrp.SurfaceDeferredRenderSession(scene, camera) 80 | # convert output tensor to PIL Image and save 81 | to_pil(rp.false_color_camera_space_normal()).save("procedural-normals.png") 82 | ``` 83 | 84 | Only 6 lines of code and we are done! The output should look like this: 85 | 86 | ```{figure} assets/procedural-normals.png 87 | :scale: 25 % 88 | :alt: Camera space normals of a procedurally-created sphere. 89 | ``` 90 | 91 | The GPU buffers can be replaced by arbitrary tensors or parameters, and the process has been fully differentiable. 92 | 93 | ```{important} 94 | The `SurfaceDeferredRenderSession` is a "session", which means you have to recreate it when you want to change the scene and/or the camera to avoid out-dated cached data. 95 | ``` 96 | 97 | ```{note} 98 | DiffRP makes heavy use of device-local caching. 99 | DiffRP operations do not support multiple GPUs on a single process. 100 | You will need to access the GPUs with a distributed paradigm, where each process only uses one GPU. 101 | ``` 102 | 103 | ### Next Steps 104 | 105 | It is recommended to go through the glTF rendering tutorial even if you do not need the functionality. It helps learning basic graphics concepts in DiffRP. 106 | 107 | ```{toctree} 108 | :maxdepth: 1 109 | gltf 110 | concepts 111 | pipeline_surface_deferred 112 | writing_a_material 113 | ptpbr 114 | pipeline_comparison 115 | ``` 116 | 117 | ## API Reference 118 | 119 | ```{note} 120 | All documented public APIs from submodules are available at the `diffrp.*` level except for resources and extension plugins. 121 | ``` 122 | 123 | ```{toctree} 124 | :maxdepth: 3 125 | generated/diffrp 126 | ``` 127 | 128 | ## Update Notes 129 | 130 | + **0.2.7**: Fix depth transform bug in surface deferred rendering pipeline. 131 | + **0.2.6**: Added `quick-render` script for previewing objects. 132 | + **0.2.5**: `pip` packaging. 133 | + **0.2.4**: Loading from trimesh scene/mesh. Support more than 16M triangles and higher than 2K resolution with corresponding upstream nvdiffrast. More complete test suite. 134 | + **0.2.3**: Make submodule contents available at top level. Make context sharing mechanism thread-safe by default. Fix `torch.load` warnings in newer torch versions. 135 | + **0.2.2**: Fixed backward compatibility due to `staticmethod`. 136 | + **0.2.1**: Add metadata functionality. 137 | + **0.2.0**: Implemented path tracing pipeline. 138 | + **0.1.5**: Fix normal and tangent attributes in static batching, refactoring to extract logic shared across rendering pipelines. 139 | + **0.1.4**: Faster GLTF loading, static batching support, minor GPU memory optimizations. 140 | + **0.1.3**: Fix instanced GLTF material loading (GPU memory explosion issue). 141 | + **0.1.2**: Fix numerical stability of coordinate transform backward. 142 | + **0.1.1**: Performance and documentation improvements. 143 | + **0.1.0**: A major rewrite of the whole package. Will try to provide backward compatibility of all documented public APIs from this version on. 144 | + **0.0.2**: Minor improvements. 145 | + **0.0.1**: First verion, fully experimental. 146 | -------------------------------------------------------------------------------- /docs/source/pipeline_comparison.md: -------------------------------------------------------------------------------- 1 | # Comparison Between Pipelines 2 | 3 | Currently, DiffRP supports two rendering pipelines: 4 | deferred lighting rasterization pipeline ({py:class}`diffrp.rendering.surface_deferred.SurfaceDeferredRenderSession`) 5 | and path tracing pipeline ({py:class}`diffrp.rendering.path_tracing.PathTracingSession`). 6 | 7 | The deferred lighting rasterization pipeline ({py:class}`diffrp.rendering.surface_deferred.SurfaceDeferredRenderSession`) supports PBR with image-based lighting, 8 | and can render various useful attributes like depth, world/camera-space normal, albedo, metallic/roughness attributes and so on. 9 | It supports analytical antialiasing and edge gradient methods as the visibility gradient. 10 | It requires ``nvdiffrast`` to be installed, which requires a working CUDA toolkit. 11 | 12 | The path tracing pipeline ({py:class}`diffrp.rendering.path_tracing.PathTracingSession`) is mainly targetting unbiased PBR. 13 | While the PBR rendering will produce albedo, world normal and first-hit positions as a by-product, 14 | rendering other attributes are not targeted by default. 15 | It is based on monte-carlo sampling of paths, so the PBR result would contain noise. 16 | You can use the denoiser, which is a differentiable version of a denoiser in OIDN, to denoise the result. 17 | It does not support common visibility gradient handling methods yet. 18 | It requires ``torchoptix`` to be installed for the best performance using RTX hardware, but can also run in pure PyTorch. 19 | -------------------------------------------------------------------------------- /docs/source/pipeline_surface_deferred.md: -------------------------------------------------------------------------------- 1 | # Surface Deferred Render Pipeline 2 | 3 | ![Surface Deferred Pipeline Diagram](assets/duffrp-deferred-rast.svg) 4 | 5 | This diagram shows the pipeline for {py:class}`diffrp.rendering.surface_deferred.SurfaceDeferredRenderSession`. 6 | 7 | Most rendering requirements can be fulfilled by only coding the material by implementing a {py:class}`diffrp.materials.base_material.SurfaceMaterial`. If more advanced render passes are wanted, you may extend the {py:class}`diffrp.rendering.surface_deferred.SurfaceDeferredRenderSession` class by inheriting from it. The dependencies shown as arrows here in the diagram is implicitly solved by the cache system, and can be requested by calling relevant functions. Of course, you can also implement custom post processing outside the render session after obtaining the results. 8 | 9 | ## Glossary 10 | 11 | ### GBuffer 12 | 13 | GBuffer is a common word for screen buffers for intermediate results in this render pipeline. 14 | 15 | Typically, we collect all information needed for the final image about objects in GBuffer, and handle lighting after we have the GBuffer. This is to decouple the complexity of objects with lights. In other words, GBuffer is a comprehensive summary of all visible objects in the screen space. 16 | 17 | GBuffer contain normals, PBR parameters (metallic, smoothness, occlusion), base color (albedo), alpha (transparency) and positions in DiffRP. They can be easily mapped into normals in various spaces (camera space, world space), NOCS, depth/distance values, which are common modalities needed for differentiable rendering projects. 18 | -------------------------------------------------------------------------------- /docs/source/ptpbr.md: -------------------------------------------------------------------------------- 1 | # Path-Traced Physically Based Rendering 2 | 3 | DiffRP implements a path tracing pipeline for high-quality physically-based rendering. 4 | 5 | We first load the [ABeautifulGame](https://github.com/KhronosGroup/glTF-Sample-Assets/tree/main/Models/ABeautifulGame) scene and add a environment light, as in the basic GLTF tutorial: 6 | 7 | ```python 8 | import torch 9 | import diffrp 10 | from diffrp.utils import * 11 | from diffrp.resources import hdris 12 | 13 | scene = diffrp.load_gltf_scene("beautygame.glb", compute_tangents=True).static_batching() 14 | scene.add_light(diffrp.ImageEnvironmentLight( 15 | intensity=1.0, color=torch.ones(3, device='cuda'), 16 | image=hdris.newport_loft().cuda(), render_skybox=True 17 | )) 18 | ``` 19 | 20 | Create a camera to look at a beautiful glazing angle: 21 | 22 | ```python 23 | camera = diffrp.PerspectiveCamera.from_orbit( 24 | h=720, w=1280, # resolution 25 | radius=0.9, azim=50, elev=10, # orbit camera position, azim/elev in degrees 26 | origin=[0, 0, 0], # orbit camera focus point 27 | fov=28.12, near=0.005, far=100.0 # intrinsics 28 | ) 29 | ``` 30 | 31 | Create the rendering session and perform rendering: 32 | 33 | ```python 34 | rp = diffrp.PathTracingSession(scene, camera, diffrp.PathTracingSessionOptions(ray_spp=8)) 35 | pbr, alpha, extras = rp.pbr() 36 | ``` 37 | 38 | Perform tone-mapping and view the image: 39 | 40 | ```python 41 | pbr_mapped = agx_base_contrast(pbr) 42 | to_pil(pbr_mapped).save("game.png") 43 | ``` 44 | 45 | ```{figure} assets/game-8spp-noisy.jpg 46 | :scale: 45 % 47 | :alt: Game noisy PBR render. 48 | ``` 49 | 50 | Due to the monte-carlo sampling there is noise in the image. 51 | You can use a higher `ray_spp` to get a less noisy image, and you can also use a denoiser. 52 | 53 | DiffRP provides a differentiable version of the HDR-ALB-NML OIDN denoiser. You can even train the denoiser if you would like! 54 | 55 | ```python 56 | denoiser = diffrp.get_denoiser() 57 | pbr_denoised = agx_base_contrast(diffrp.run_denoiser(denoiser, pbr, linear_to_srgb(extras['albedo']), extras['world_normal'])) 58 | to_pil(pbr_denoised).save("gamed.png") 59 | ``` 60 | 61 | ```{figure} assets/game-1024spp-denoised.jpg 62 | :scale: 45 % 63 | :alt: Game denoised PBR render. 64 | ``` 65 | 66 | Transparent film is not well supported for now, `render_skybox` will not affect the background region in the radiance output. 67 | You can still combine with the `alpha` output to obtain a transparent film image, but you may see some background bleeding into the edges. 68 | 69 | ```python 70 | to_pil(float4(pbr_denoised, alpha)).save("gameda.png") 71 | ``` 72 | 73 | ```{figure} assets/game-8spp-denoised-alpha.jpg 74 | :scale: 45 % 75 | :alt: Game denoised PBR render with transparent film. 76 | ``` 77 | 78 | Complete example: 79 | 80 | ```python 81 | import torch 82 | import diffrp 83 | from diffrp.utils import * 84 | from diffrp.resources import hdris 85 | 86 | torch.set_grad_enabled(False) 87 | scene = diffrp.load_gltf_scene("beautygame.glb", True).static_batching() 88 | scene.add_light(diffrp.ImageEnvironmentLight( 89 | intensity=1.0, color=torch.ones(3, device='cuda'), 90 | image=hdris.newport_loft().cuda(), render_skybox=True 91 | )) 92 | camera = diffrp.PerspectiveCamera.from_orbit( 93 | h=720, w=1280, # resolution 94 | radius=0.9, azim=50, elev=10, # orbit camera position, azim/elev in degrees 95 | origin=[0, 0, 0], # orbit camera focus point 96 | fov=28.12, near=0.005, far=100.0 # intrinsics 97 | ) 98 | rp = diffrp.PathTracingSession(scene, camera, diffrp.PathTracingSessionOptions(ray_spp=8)) 99 | pbr, alpha, extras = rp.pbr() 100 | denoiser = diffrp.get_denoiser() 101 | pbr_mapped = agx_base_contrast(pbr) 102 | pbr_denoised = agx_base_contrast(diffrp.run_denoiser(denoiser, pbr, linear_to_srgb(extras['albedo']), extras['world_normal'])) 103 | to_pil(pbr_mapped).save("game.png") 104 | to_pil(pbr_denoised).save("gamed.png") 105 | to_pil(float4(pbr_denoised, alpha)).save("gameda.png") 106 | ``` 107 | -------------------------------------------------------------------------------- /docs/source/templates/module.rst_t: -------------------------------------------------------------------------------- 1 | {{ basename | e | heading }} 2 | 3 | .. automodule:: {{ qualname }} 4 | {%- for option in automodule_options %} 5 | :{{ option }}: 6 | {%- endfor %} 7 | -------------------------------------------------------------------------------- /docs/source/templates/package.rst_t: -------------------------------------------------------------------------------- 1 | {{ pkgname | e | heading }} 2 | .. automodule:: {{ pkgname }} 3 | {%- for option in automodule_options %} 4 | :{{ option }}: 5 | {%- endfor %} 6 | {% if subpackages or submodules %} 7 | .. toctree:: 8 | {% for docname in (subpackages or []) + (submodules or []) %} 9 | {{ docname }} 10 | {%- endfor %} 11 | {% endif %} 12 | -------------------------------------------------------------------------------- /docs/source/writing_a_material.md: -------------------------------------------------------------------------------- 1 | # Writing a Material 2 | 3 | Most rendering requirements can be fulfilled by only coding the material by implementing a {py:class}`diffrp.materials.base_material.SurfaceMaterial`. 4 | 5 | DiffRP is designed with modularity and extendability in mind. All interfaces are in the simple form of PyTorch function calls. For the default interpolator implementations, tensors you see in `SurfaceInput` will have shape (B, C) where B is batch number of pixels and C is the relevant channels, for example, positions and normals have 3 channels while RGBA colors have 4. 6 | 7 | We show two examples: procedural effects and neural networks. 8 | 9 | ## Procedural Effects 10 | 11 | We implement a material where the RGB is defined as functions of the world space coordinates. 12 | For points with `z > 0.5`, we additionally increase the values by a star-like shape. 13 | 14 | Note that within DiffRP you can use the GLSL-like attributes to access slices in Tensors. 15 | 16 | ```python 17 | import torch 18 | import diffrp 19 | from diffrp.utils import * 20 | 21 | 22 | class ProceduralMaterial(diffrp.SurfaceMaterial): 23 | 24 | def shade(self, su: diffrp.SurfaceUniform, si: diffrp.SurfaceInput) -> diffrp.SurfaceOutputStandard: 25 | p = si.world_pos 26 | r = torch.exp(torch.sin(p.x) + torch.cos(p.y) - 2) 27 | g = torch.exp(torch.sin(p.y) + torch.cos(p.z) - 2) 28 | b = torch.exp(torch.sin(p.z) + torch.cos(p.x) - 2) 29 | v = torch.where(p.z > 0.5, 0.5 * torch.exp(-p.x * p.x * p.y * p.y * 1000), 0) 30 | albedo = float3(r, g, b) + v 31 | return diffrp.SurfaceOutputStandard(albedo=albedo) 32 | ``` 33 | 34 | We can use the quick-start example for visualization: 35 | 36 | ```python 37 | import diffrp 38 | import trimesh.creation 39 | from diffrp.utils import * 40 | 41 | mesh = trimesh.creation.icosphere(radius=0.8) 42 | scene = diffrp.Scene().add_mesh_object(diffrp.MeshObject(ProceduralMaterial(), gpu_f32(mesh.vertices), gpu_i32(mesh.faces))) 43 | camera = diffrp.PerspectiveCamera(h=2048, w=2048) 44 | rp = diffrp.SurfaceDeferredRenderSession(scene, camera) 45 | to_pil(rp.albedo_srgb()).save("procedural-albedo.png") 46 | ``` 47 | 48 | Note that we convert the values from linear to sRGB for better visuals. 49 | 50 | ```{figure} assets/procedural-albedo.png 51 | :scale: 50 % 52 | :alt: Procedural albedo material. 53 | ``` 54 | 55 | ## Neural Networks 56 | 57 | As the input and outputs are expect to be a batch of features, and they are already in native PyTorch tensor interfaces, it is quite easy to integrate a neural network in DiffRP. It is just like calling the other functions. 58 | 59 | In this `NeuralMaterial`, we create a neural network and evaluate it for albedo output. 60 | 61 | ```python 62 | class NeuralMaterial(diffrp.SurfaceMaterial): 63 | def __init__(self) -> None: 64 | super().__init__() 65 | self.net = torch.nn.Sequential( 66 | torch.nn.Linear(3, 16), 67 | torch.nn.SiLU(), 68 | torch.nn.Linear(16, 16), 69 | torch.nn.SiLU(), 70 | torch.nn.Linear(16, 3), 71 | torch.nn.Sigmoid() 72 | ).cuda() 73 | 74 | def shade(self, su: diffrp.SurfaceUniform, si: diffrp.SurfaceInput) -> diffrp.SurfaceOutputStandard: 75 | return diffrp.SurfaceOutputStandard(albedo=self.net(si.world_normal)) 76 | ``` 77 | 78 | We use the same testing code: 79 | 80 | ```python 81 | mesh = trimesh.creation.icosphere(radius=0.8) 82 | material = NeuralMaterial() 83 | scene = diffrp.Scene().add_mesh_object(diffrp.MeshObject(material, gpu_f32(mesh.vertices), gpu_i32(mesh.faces))) 84 | camera = diffrp.PerspectiveCamera(h=512, w=512) 85 | rp = diffrp.SurfaceDeferredRenderSession(scene, camera) 86 | to_pil(rp.albedo_srgb()).save("neural-albedo.png") 87 | ``` 88 | 89 | ```{figure} assets/neural-albedo.png 90 | :scale: 50 % 91 | :alt: Neural albedo material. 92 | ``` 93 | 94 | You can also back-propagate the result -- it just works. 95 | 96 | ```python 97 | pred = rp.albedo_srgb() 98 | torch.nn.functional.mse_loss(pred, torch.rand_like(pred)).backward() 99 | print(material.net[0].weight.grad) 100 | ``` 101 | 102 | Output: 103 | 104 | ``` 105 | tensor([[ 4.9758e-05, 4.8679e-05, 3.0747e-04], 106 | [-1.9968e-05, 1.7120e-06, 3.6997e-04], 107 | [-1.5356e-05, 1.8661e-05, 1.6705e-04], 108 | [-2.2664e-05, -1.2767e-05, -7.8220e-05], 109 | [ 8.0282e-06, 1.1283e-05, -1.3932e-04], 110 | [-3.6454e-06, 4.1644e-05, 2.5145e-04], 111 | [-2.8950e-06, -8.8152e-06, -3.0734e-04], 112 | [ 2.8373e-05, 2.3682e-05, 6.7197e-04], 113 | [-2.3877e-05, 5.6425e-05, 5.7738e-04], 114 | [ 4.8350e-05, -9.0493e-05, -1.0172e-03], 115 | [-1.4312e-05, -4.0098e-05, 2.5461e-04], 116 | [ 7.0399e-05, 2.5832e-05, 4.8466e-04], 117 | [-3.7648e-05, -2.5028e-05, -1.4115e-04], 118 | [ 7.4847e-05, 1.7914e-05, 5.5394e-04], 119 | [ 3.1179e-06, -4.5348e-06, 1.2603e-05], 120 | [ 7.7181e-05, 3.2846e-05, -7.4523e-04]], device='cuda:0') 121 | ``` 122 | 123 | In most cases, neural networks are agnostic to color spaces, so it will make little difference if we assume it is in sRGB or linear space (by converting or not converting the color space at output). 124 | 125 | You may also apply losses on raw outputs from this neural network: 126 | 127 | ```python 128 | pred = rp.albedo() 129 | torch.nn.functional.mse_loss(pred, torch.rand_like(pred)).backward() 130 | ``` 131 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | 4 | with open("README.md", "rb") as fh: 5 | long_description = fh.read().decode() 6 | with open("diffrp/version.py", "r") as fh: 7 | exec(fh.read()) 8 | __version__: str 9 | 10 | 11 | def packages(): 12 | return setuptools.find_packages(include=['diffr*']) 13 | 14 | 15 | setuptools.setup( 16 | name="diffrp", 17 | version=__version__, 18 | author="flandre.info", 19 | author_email="flandre@scarletx.cn", 20 | description="Differentiable Render Pipelines with PyTorch.", 21 | long_description=long_description, 22 | long_description_content_type="text/markdown", 23 | url="https://github.com/eliphatfs/diffrp", 24 | packages=packages(), 25 | classifiers=[ 26 | "Programming Language :: Python :: 3 :: Only", 27 | "License :: OSI Approved :: Apache Software License", 28 | "Operating System :: OS Independent", 29 | ], 30 | install_requires=[ 31 | 'rich', 32 | 'numpy', 33 | 'torch', 34 | 'pyexr', 35 | 'calibur', 36 | 'trimesh[easy]', 37 | 'torch_redstone', 38 | 'typing_extensions', 39 | 'diffrp-nvdiffrast', 40 | ], 41 | package_data={"diffrp.plugins": ["**/*.c", "**/*.h"], "diffrp": ["resources/**/*.*"]}, 42 | exclude_package_data={ 43 | 'diffrp': ["*.pyc"], 44 | }, 45 | entry_points=dict( 46 | console_scripts=[ 47 | "diffrp=diffrp.scripts.main:main" 48 | ] 49 | ), 50 | python_requires='~=3.7' 51 | ) 52 | --------------------------------------------------------------------------------