Upload 65 files

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+asset_visualization/armor.gif filter=lfs diff=lfs merge=lfs -text
+meshes_target/jacket_sdf_new.obj filter=lfs diff=lfs merge=lfs -text
+meshes/UV.png filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,88 @@
+# Use CUDA base image
+FROM nvidia/cuda:11.8-devel-ubuntu20.04
+
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    python3.8 \
+    python3.8-dev \
+    python3-pip \
+    git \
+    wget \
+    curl \
+    build-essential \
+    cmake \
+    ninja-build \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    libsm6 \
+    libxext6 \
+    libxrender-dev \
+    libgomp1 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create symbolic link for python
+RUN ln -s /usr/bin/python3.8 /usr/bin/python
+
+# Upgrade pip
+RUN python -m pip install --upgrade pip
+
+# Set working directory
+WORKDIR /app
+
+# Copy requirements first for better caching
+COPY requirements.txt .
+
+# Install PyTorch with CUDA support
+RUN pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu118
+
+# Install basic dependencies
+RUN pip install -r requirements.txt
+
+# Install CLIP
+RUN pip install git+https://github.com/openai/CLIP.git
+
+# Create packages directory
+RUN mkdir -p packages
+
+# Install nvdiffrast
+WORKDIR /app/packages
+RUN git clone https://github.com/NVlabs/nvdiffrast.git && \
+    cd nvdiffrast && \
+    pip install .
+
+# Install PyTorch3D
+WORKDIR /app/packages
+RUN git clone https://github.com/facebookresearch/pytorch3d.git && \
+    cd pytorch3d && \
+    pip install .
+
+# Install Fashion-CLIP
+WORKDIR /app/packages
+RUN git clone https://github.com/patrickjohncyh/fashion-clip.git && \
+    cd fashion-clip && \
+    pip install appdirs boto3 annoy validators transformers datasets
+
+# Return to app directory
+WORKDIR /app
+
+# Copy application files
+COPY . .
+
+# Create necessary directories
+RUN mkdir -p outputs meshes meshes_target
+
+# Make setup script executable
+RUN chmod +x setup_spaces.py
+
+# Expose port
+EXPOSE 7860
+
+# Set the default command
+CMD ["python", "app.py"] 

LICENSE.md

@@ -0,0 +1,159 @@
+# Attribution-NonCommercial 4.0 International
+
+> *Creative Commons Corporation (“Creative Commons”) is not a law firm and does not provide legal services or legal advice. Distribution of Creative Commons public licenses does not create a lawyer-client or other relationship. Creative Commons makes its licenses and related information available on an “as-is” basis. Creative Commons gives no warranties regarding its licenses, any material licensed under their terms and conditions, or any related information. Creative Commons disclaims all liability for damages resulting from their use to the fullest extent possible.*
+>
+> ### Using Creative Commons Public Licenses
+>
+> Creative Commons public licenses provide a standard set of terms and conditions that creators and other rights holders may use to share original works of authorship and other material subject to copyright and certain other rights specified in the public license below. The following considerations are for informational purposes only, are not exhaustive, and do not form part of our licenses.
+>
+> * __Considerations for licensors:__ Our public licenses are intended for use by those authorized to give the public permission to use material in ways otherwise restricted by copyright and certain other rights. Our licenses are irrevocable. Licensors should read and understand the terms and conditions of the license they choose before applying it. Licensors should also secure all rights necessary before applying our licenses so that the public can reuse the material as expected. Licensors should clearly mark any material not subject to the license. This includes other CC-licensed material, or material used under an exception or limitation to copyright. [More considerations for licensors](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensors).
+>
+> * __Considerations for the public:__ By using one of our public licenses, a licensor grants the public permission to use the licensed material under specified terms and conditions. If the licensor’s permission is not necessary for any reason–for example, because of any applicable exception or limitation to copyright–then that use is not regulated by the license. Our licenses grant only permissions under copyright and certain other rights that a licensor has authority to grant. Use of the licensed material may still be restricted for other reasons, including because others have copyright or other rights in the material. A licensor may make special requests, such as asking that all changes be marked or described. Although not required by our licenses, you are encouraged to respect those requests where reasonable. [More considerations for the public](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensees).
+
+## Creative Commons Attribution-NonCommercial 4.0 International Public License
+
+By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions.
+
+### Section 1 – Definitions.
+
+a. __Adapted Material__ means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image.
+
+b. __Adapter's License__ means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License.
+
+c. __Copyright and Similar Rights__ means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights.
+
+d. __Effective Technological Measures__ means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements.
+
+e. __Exceptions and Limitations__ means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material.
+
+f. __Licensed Material__ means the artistic or literary work, database, or other material to which the Licensor applied this Public License.
+
+g. __Licensed Rights__ means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license.
+
+h. __Licensor__ means the individual(s) or entity(ies) granting rights under this Public License.
+
+i. __NonCommercial__ means not primarily intended for or directed towards commercial advantage or monetary compensation. For purposes of this Public License, the exchange of the Licensed Material for other material subject to Copyright and Similar Rights by digital file-sharing or similar means is NonCommercial provided there is no payment of monetary compensation in connection with the exchange.
+
+j. __Share__ means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them.
+
+k. __Sui Generis Database Rights__ means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world.
+
+l. __You__ means the individual or entity exercising the Licensed Rights under this Public License. Your has a corresponding meaning.
+
+### Section 2 – Scope.
+
+a. ___License grant.___
+
+   1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to:
+
+       A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and
+
+       B. produce, reproduce, and Share Adapted Material for NonCommercial purposes only.
+
+   2. __Exceptions and Limitations.__ For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions.
+
+   3. __Term.__ The term of this Public License is specified in Section 6(a).
+
+   4. __Media and formats; technical modifications allowed.__ The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material.
+
+   5. __Downstream recipients.__
+
+        A. __Offer from the Licensor – Licensed Material.__ Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License.
+
+        B. __No downstream restrictions.__ You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material.
+
+   6. __No endorsement.__ Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i).
+
+b. ___Other rights.___
+
+   1. Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise.
+
+   2. Patent and trademark rights are not licensed under this Public License.
+
+   3. To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties, including when the Licensed Material is used other than for NonCommercial purposes.
+
+### Section 3 – License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the following conditions.
+
+a. ___Attribution.___
+
+   1. If You Share the Licensed Material (including in modified form), You must:
+
+       A. retain the following if it is supplied by the Licensor with the Licensed Material:
+
+         i. identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated);
+
+         ii. a copyright notice;
+
+         iii. a notice that refers to this Public License;
+
+         iv. a notice that refers to the disclaimer of warranties;
+
+         v. a URI or hyperlink to the Licensed Material to the extent reasonably practicable;
+
+       B. indicate if You modified the Licensed Material and retain an indication of any previous modifications; and
+
+       C. indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License.
+
+   2. You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information.
+
+   3. If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable.
+
+   4. If You Share Adapted Material You produce, the Adapter's License You apply must not prevent recipients of the Adapted Material from complying with this Public License.
+
+### Section 4 – Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material:
+
+a. for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database for NonCommercial purposes only;
+
+b. if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material; and
+
+c. You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights.
+
+### Section 5 – Disclaimer of Warranties and Limitation of Liability.
+
+a. __Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You.__
+
+b. __To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You.__
+
+c. The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability.
+
+### Section 6 – Term and Termination.
+
+a. This Public License applies for the term of the Copyright and Similar Rights licensed here. However, if You fail to comply with this Public License, then Your rights under this Public License terminate automatically.
+
+b. Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates:
+
+   1. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or
+
+   2. upon express reinstatement by the Licensor.
+
+   For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License.
+
+c. For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License.
+
+d. Sections 1, 5, 6, 7, and 8 survive termination of this Public License.
+
+### Section 7 – Other Terms and Conditions.
+
+a. The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed.
+
+b. Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License.
+
+### Section 8 – Interpretation.
+
+a. For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License.
+
+b. To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions.
+
+c. No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor.
+
+d. Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority.
+
+> Creative Commons is not a party to its public licenses. Notwithstanding, Creative Commons may elect to apply one of its public licenses to material it publishes and in those instances will be considered the “Licensor.” Except for the limited purpose of indicating that material is shared under a Creative Commons public license or as otherwise permitted by the Creative Commons policies published at [creativecommons.org/policies](http://creativecommons.org/policies), Creative Commons does not authorize the use of the trademark “Creative Commons” or any other trademark or logo of Creative Commons without its prior written consent including, without limitation, in connection with any unauthorized modifications to any of its public licenses or any other arrangements, understandings, or agreements concerning use of licensed material. For the avoidance of doubt, this paragraph does not form part of the public licenses.
+>
+> Creative Commons may be contacted at creativecommons.org

NeuralJacobianFields/MeshProcessor.py

@@ -0,0 +1,473 @@
+from multiprocessing import process
+import warnings
+warnings.filterwarnings("ignore")
+
+
+from scipy.sparse import load_npz, save_npz
+from .PoissonSystem import poisson_system_matrices_from_mesh, PoissonSystemMatrices, SparseMat
+import os
+import trimesh
+from easydict import EasyDict
+import numpy
+import numpy as np
+import scipy
+import scipy.sparse
+import igl
+from scipy.sparse import save_npz
+from time import time
+import torch
+
+NUM_SAMPLES = 1024
+WKS_DIM = 100
+
+class MeshProcessor:
+    '''
+    Extracts all preprocessing-related data (sample points  for pointnet; wave-kernel-signature, etc.)
+    '''
+    def __init__(self, vertices, faces, ttype, source_dir=None,from_file = False,
+                 cpuonly=False, load_wks_samples=False, load_wks_centroids=False,
+                compute_splu=True, load_splu=False):
+        '''
+        :param vertices:
+        :param faces:
+        :param ttype: the torch data type to use (float, half, double)
+        :param source_dir: the directory to load the preprocessed data from; if given, will try to load the data before computing, if not given, always compute
+        '''
+        
+        self.ttype = ttype
+        self.num_samples = NUM_SAMPLES
+        self.vertices = vertices.squeeze()
+        self.faces = faces.squeeze()
+        self.normals =  igl.per_vertex_normals(self.vertices, self.faces)
+        # self.__use_wks = use_wks
+        self.samples = EasyDict()
+        self.samples.xyz = None
+        self.samples.normals = None
+        self.samples.wks = None
+        self.centroids = EasyDict()
+        self.centroids.points_and_normals = None
+        self.centroids.wks = None
+        self.diff_ops = EasyDict()
+        self.diff_ops.splu = EasyDict()
+        self.diff_ops.splu.L = None
+        self.diff_ops.splu.U = None
+        self.diff_ops.splu.perm_c = None
+        self.diff_ops.splu.perm_r = None
+        self.diff_ops.frames = None
+        self.diff_ops.rhs = None
+        self.diff_ops.grad = None
+        self.diff_ops.poisson_sys_mat = None
+        self.faces_wks = None
+        self.vert_wks = None
+        self.diff_ops.poisson = None
+        self.source_dir = source_dir
+        self.from_file = from_file
+        self.cpuonly = cpuonly
+        self.load_wks_samples = load_wks_samples
+        self.load_wks_centroids = load_wks_centroids
+        self.compute_splu = compute_splu
+        self.load_splu = load_splu
+
+    @staticmethod
+    def meshprocessor_from_directory(source_dir, ttype, cpuonly=False, load_wks_samples=False, load_wks_centroids=False):
+        try:
+            vertices = np.load(os.path.join(source_dir, "vertices.npy"))
+            faces = np.load(os.path.join(source_dir, "faces.npy"))
+        except:
+            print(os.path.join(source_dir, "vertices.npy"))
+            import traceback
+            traceback.print_exc()
+        return MeshProcessor(vertices,faces,ttype,source_dir, cpuonly=cpuonly, load_wks_samples=load_wks_samples, load_wks_centroids=load_wks_centroids, compute_splu=False)
+
+    @staticmethod
+    def meshprocessor_from_file(fname, ttype, cpuonly=False, load_wks_samples=False, load_wks_centroids=False):
+        if fname[-4:] == '.obj':
+            V, _, _, F, _, _ = igl.read_obj(fname)
+        elif fname[-4:] == '.off':
+            V,F,_ = igl.read_off(fname)
+        elif fname[-4:] == '.ply':
+            V,F = igl.read_triangle_mesh(fname)
+        return MeshProcessor(V,F,ttype,os.path.dirname(fname),True, cpuonly=cpuonly, load_wks_samples=load_wks_samples, load_wks_centroids=load_wks_centroids, compute_splu=False)
+
+    @staticmethod
+    def meshprocessor_from_array(vertices, faces, source_dir, ttype, cpuonly=False, load_wks_samples=False, load_wks_centroids=False):
+        return MeshProcessor(vertices,faces,ttype,source_dir, cpuonly=cpuonly, load_wks_samples=load_wks_samples, load_wks_centroids=load_wks_centroids, compute_splu=False)
+
+
+
+    def get_vertices(self):
+        return self.vertices
+
+    def get_faces(self):
+        return self.faces
+
+    def load_centroids(self):
+        self.centroids.points_and_normals = np.load(os.path.join(self.source_dir, "centroids_and_normals.npy"))
+        if self.load_wks_centroids:
+            self.centroids.wks = np.load(os.path.join(self.source_dir, "centroids_wks.npy"))
+
+    def get_samples(self):
+        if self.samples.xyz is None:
+            if True:
+                try:
+                    self.load_samples()
+                except Exception as e:
+                    self.compute_samples()
+                    self.save_samples()
+        return self.samples
+
+    def load_samples(self):
+        if self.samples.xyz is None:
+            self.samples.xyz = np.load(os.path.join(self.source_dir, 'samples.npy'))
+        if self.samples.normals is None:
+            self.samples.normals = np.load(os.path.join(self.source_dir, 'samples_normals.npy'))
+        if self.load_wks_samples:
+            if  self.samples.wks is None:
+                self.samples.wks = np.load(os.path.join(self.source_dir, 'samples_wks.npy'))
+            if self.centroids.wks is None:
+                self.centroids.wks = np.load(os.path.join(self.source_dir, 'centroid_wks.npy'))
+
+    def save_samples(self):
+        os.makedirs(self.source_dir, exist_ok=True)
+        np.save(os.path.join(self.source_dir, 'samples.npy'), self.samples.xyz)
+        np.save(os.path.join(self.source_dir, 'samples_normals.npy'), self.samples.normals)
+        if self.load_wks_samples:
+            np.save(os.path.join(self.source_dir, 'samples_wks.npy'), self.samples.wks)
+            np.save(os.path.join(self.source_dir, 'centroid_wks.npy'), self.centroids.wks)
+
+    def compute_samples(self):
+        sstime = time()
+        if self.load_wks_centroids or self.load_wks_centroids:
+            self.computeWKS()
+        # print(f"WKS {time() - sstime}")
+        pt_samples, normals_samples, wks_samples, bary = self.sample_points( self.num_samples)
+        self.samples.xyz = pt_samples
+        self.samples.normals = normals_samples
+        self.samples.wks = wks_samples
+        self.centroids.wks = self.faces_wks
+
+    def get_centroids(self):
+        if self.centroids.points_and_normals is None:
+            if True:
+                try:
+                    self.load_centroids()
+                except Exception as e:
+                    self.compute_centroids()
+                    # self.save_centroids() # centroid WKS and samples WKS are intertwined right now and you cannot really use one without the other. So this is redondont with function save_samples
+        return self.centroids
+
+    def compute_centroids(self):
+        m = trimesh.Trimesh(vertices=self.vertices, faces=self.faces, process=False)
+        self.centroids.points_and_normals = np.hstack((np.mean(m.triangles, axis=1), m.face_normals))
+        self.get_samples()# this is to compute WKS for centroids
+
+    def get_differential_operators(self):
+        if self.diff_ops.grad is None:
+            if True:
+                try:
+                    self.load_differential_operators()
+                except Exception as e:
+                    warnings.warn(f'while loading data, got file not exists exception: {e} ')
+                    self.compute_differential_operators()
+                    self.save_differential_operators()
+        if self.load_splu:
+            self.get_poisson_system()
+        return self.diff_ops
+
+
+    def load_poisson_system(self):
+        try:
+            self.diff_ops.splu.L = load_npz(os.path.join(self.source_dir, 'lap_L.npz'))
+            self.diff_ops.splu.U = load_npz(os.path.join(self.source_dir, 'lap_U.npz'))
+            self.diff_ops.splu.perm_c = np.load(os.path.join(self.source_dir, 'lap_perm_c.npy'))
+            self.diff_ops.splu.perm_r = np.load(os.path.join(self.source_dir, 'lap_perm_r.npy'))
+        except:
+            print(f"FAILED load poisson on: {os.path.join(self.source_dir)}")
+            raise Exception("FAILED load poisson on: {os.path.join(self.source_dir)}")
+
+
+    def load_differential_operators(self): 
+        self.diff_ops.rhs = SparseMat.from_coo(load_npz(os.path.join(self.source_dir, 'new_rhs.npz')), ttype=torch.float64)
+        self.diff_ops.grad = SparseMat.from_coo(load_npz(os.path.join(self.source_dir, 'new_grad.npz')), ttype=torch.float64)
+        self.diff_ops.frames = np.load(os.path.join(self.source_dir, 'w.npy'))
+        self.diff_ops.laplacian = SparseMat.from_coo(load_npz(os.path.join(self.source_dir, 'laplacian.npz')), ttype=torch.float64)
+
+    def save_differential_operators(self):  
+        save_npz(os.path.join(self.source_dir, 'new_rhs.npz'), self.diff_ops.rhs.to_coo())
+        save_npz(os.path.join(self.source_dir, 'new_grad.npz'), self.diff_ops.grad.to_coo())
+        np.save(os.path.join(self.source_dir, 'w.npy'), self.diff_ops.frames)
+        save_npz(os.path.join(self.source_dir, 'laplacian.npz'), self.diff_ops.laplacian.to_coo())
+
+    def compute_differential_operators(self):
+        '''
+        process the given mesh
+        '''
+        poisson_sys_mat = poisson_system_matrices_from_mesh(V= self.vertices, F=self.faces,  cpuonly=self.cpuonly)
+        self.diff_ops.grad = poisson_sys_mat.igl_grad
+        self.diff_ops.rhs = poisson_sys_mat.rhs
+        self.diff_ops.laplacian = poisson_sys_mat.lap
+        self.diff_ops.frames = poisson_sys_mat.w
+        self.diff_ops.poisson_sys_mat = poisson_sys_mat
+    
+
+    def compute_poisson(self):
+        poissonsolver = poissonbuilder.compute_poisson_solver_from_laplacian(compute_splu=self.compute_splu)
+        # new_grad = poissonbuilder.get_new_grad() # This is now done in poisson_system_matrices_from_mesh
+        if self.compute_splu:
+            self.diff_ops.splu.L, self.diff_ops.splu.U , self.diff_ops.splu.perm_c , self.diff_ops.splu.perm_r = poissonbuilder.compute_splu()
+        self.diff_ops.frames = poissonbuilder.w
+
+
+    def prepare_differential_operators_for_use(self,ttype):
+        diff_ops = self.get_differential_operators() # call 1
+        ## WARNING : we commented these two lines because they seemed redundant.
+        if self.diff_ops.poisson_sys_mat is None: # not created if loaded from disk the diff ops
+            diff_ops.poisson_sys_mat = PoissonSystemMatrices(self.vertices, self.faces, diff_ops.grad, diff_ops.rhs, diff_ops.frames, ttype, lap = diff_ops.laplacian,  cpuonly=self.cpuonly)
+        
+        self.diff_ops.poisson_solver = diff_ops.poisson_sys_mat.create_poisson_solver() # call 2
+        self.diff_ops.MyCuSPLU_solver = diff_ops.poisson_sys_mat.create_poisson_solver() #create_poisson_solver_from_splu_old
+
+    def get_writeable(self):
+        '''
+        get dictionaries to write numpy and npz
+        :return: two args, np, npz, each dicts with field_name --> data to save
+        '''
+        out_np = {}
+        out_npz = {}
+        out_np['vertices'] = self.vertices
+        out_np['faces'] = self.faces
+        if self.samples is not None:
+            out_np["samples"] = self.samples.xyz
+            out_np["samples_normals"] = self.samples.normals
+            out_np["samples_wks"] = self.samples.wks
+        if self.centroids is not None:
+            out_np["centroids_wks"] = self.centroids.wks
+            out_np["centroids_and_normals"] = self.centroids.points_and_normals
+        if self.diff_ops is not None:
+            out_np['lap_perm_c'] = self.diff_ops.splu.perm_c
+            out_np['lap_perm_r'] = self.diff_ops.splu.perm_r
+            out_np['w'] =self.diff_ops.frames
+            out_npz['new_grad'] = self.diff_ops.grad.to_coo()
+            out_npz['new_rhs'] = self.diff_ops.rhs
+            out_npz['lap_L'] = self.diff_ops.splu.L
+            out_npz['lap_U'] = self.diff_ops.splu.U
+            out_npz['lap'] = self.diff_ops.poisson.lap
+        return {key: value for key, value in out_np.items() if value is not None}, {key: value for key, value in out_npz.items() if value is not None}
+        
+    def get_data(self, key,file_type = 'npy'):
+        if key == 'samples':
+            return self.get_samples().xyz
+        elif key == "samples_normals":
+            return self.get_samples().normals
+        elif key == "samples_wks":
+            return self.get_samples().wks
+        elif key == 'vertices':
+            return self.vertices
+        elif key == 'faces':
+            return self.faces
+        if file_type == 'npy':
+            return np.load(os.path.join(self.source_dir, f'{key}.npy'))
+        elif file_type == 'npz':
+            return load_npz(os.path.join(self.source_dir, f'{key}.npz'))
+        else:
+            raise RuntimeError("wrong file type")
+
+    def computeWKS(self):
+        if self.faces_wks is  None or self.vert_wks is  None:
+            st = time()
+            w = WaveKernelSignature(self.vertices, self.faces, top_k_eig=50)
+            w.compute()
+            print(f"Ellapsed {time() - st}")
+            wk = w.wks
+            faces_wks = np.zeros((self.faces.shape[0], wk.shape[1]))
+            for i in range(3):
+                faces_wks += wk[self.faces[:, i], :]
+            faces_wks /= 3
+            self.faces_wks = faces_wks
+            self.vert_wks = wk
+            assert (self.faces_wks.shape[0] == self.faces.shape[0])
+            assert (self.vert_wks.shape[0] == self.vertices.shape[0])
+
+
+    def sample_points(self, n):
+        bary, found_faces = igl.random_points_on_mesh(n, self.vertices, self.faces)
+        vert_ind =  self.faces[found_faces]
+        point_samples =  self.vertices[vert_ind[:,0]] * bary[:,0:1] + self.vertices[vert_ind[:,1]] * bary[:,1:2] + self.vertices[vert_ind[:,2]] * bary[:,2:3]
+        normal_samples = self.normals[vert_ind[:,0]] * bary[:,0:1] + self.normals[vert_ind[:,1]] * bary[:,1:2] + self.normals[vert_ind[:,2]] * bary[:,2:3]
+        wks_samples = None
+        if self.load_wks_centroids or self.load_wks_samples:
+            wks_samples = self.vert_wks[vert_ind[:,0]] * bary[:,0:1] + self.vert_wks[vert_ind[:,1]] * bary[:,1:2] + self.vert_wks[vert_ind[:,2]] * bary[:,2:3]
+        return point_samples, normal_samples, wks_samples, bary
+
+# This is insane to me
+def sample_points(V, F, n):
+    '''
+    samples n points on the given mesh, along with normals and wks. Also return WKS of original faces (by averaging wks of 3 vertices of each face)
+    :return:
+    '''
+    newF = F
+    newV = V
+    for iter in range(n):
+        newV, newF = _sample_point(newV, newF)
+
+    w = WaveKernelSignature(newV, newF, top_k_eig=100)
+    w.compute()
+    wk = w.wks
+    sample_ks = wk[len(V):, :]
+    org_ks = wk[:len(V), :]
+    normals = igl.per_vertex_normals(newV, newF)
+    normals = normals[len(V):, :]
+
+    # get per-face wks by averaging its vertices
+    faces_wks = np.zeros((F.shape[0], org_ks.shape[1]))
+    for i in range(3):
+        faces_wks += org_ks[F[:, i], :]
+    faces_wks /= 3
+    return newV[len(V):, :], normals, sample_ks, faces_wks
+
+
+def _sample_point(VV, FF):
+    while (True):
+        bary, found_faces = igl.random_points_on_mesh(1, VV, FF)
+        if (found_faces >= FF.shape[0]):
+            continue
+        # use to be 0.01
+        if not numpy.any(bary < 0.05):
+            break
+    ret = numpy.zeros((1, VV.shape[1]))
+    for i in range(VV.shape[1]):
+        res = np.multiply(VV[FF[found_faces, :], i], bary)
+        ret[:, i] = np.sum(res)
+    newF = FF
+    new_index = len(VV)
+    new_tris = _insert_triangle(FF[found_faces, :], new_index)
+    newF = numpy.concatenate((newF, new_tris), axis=0)
+    newF = numpy.delete(newF, found_faces, axis=0)
+    newV = numpy.concatenate((VV, ret), 0)
+    return newV, newF
+
+
+def _insert_triangle(old_tri, new_index):
+    d = new_index
+    a, b, c = (old_tri[0], old_tri[1], old_tri[2])
+    new_tris = numpy.array([[a, b, d], [b, c, d], [c, a, d]])
+    return new_tris
+
+
+
+
+
+class WaveKernelSignatureError(Exception):
+    pass
+
+class WaveKernelSignature:
+    '''
+    Computes wave kernel signature for a given mesh
+    '''
+    def __init__(self,
+                 vertices,
+                 faces,
+                 top_k_eig=200,
+                 timestamps=WKS_DIM):
+        # vertices, faces are both numpy arrays.
+        self.vertices = vertices
+        self.faces = faces
+
+        # self.vertices_gpu = torch.from_numpy(vertices).cuda()
+        # self.faces_gpu = torch.from_numpy(faces).cuda()
+
+        self.top_k_eig = top_k_eig
+        self.timestamps = timestamps
+        self.max_iter = 10000
+
+    def compute(self):
+        '''
+        compute the wks. Afterwards WKS stores in self.wks
+        '''
+        cp = igl.connected_components(igl.adjacency_matrix(self.faces))
+        assert(cp[0]==1), f"{cp}"
+        L = -igl.cotmatrix(self.vertices, self.faces) # this is fast 0.04 seconds
+        M = igl.massmatrix(self.vertices, self.faces, igl.MASSMATRIX_TYPE_VORONOI)
+        # assert(not numpy.any(numpy.isinf(L)))
+        try:
+            try:
+                self.eig_vals, self.eig_vecs = scipy.sparse.linalg.eigsh(
+                    L, self.top_k_eig, M, sigma=0, which='LM', maxiter=self.max_iter)
+            except:
+                self.eig_vals, self.eig_vecs = scipy.sparse.linalg.eigsh(
+                    L, self.top_k_eig, M, sigma=1e-4, which='LM', maxiter=self.max_iter)                
+        except:
+            raise WaveKernelSignatureError("Error in computing WKS")
+
+        # print(np.linalg.norm(self.eig_vecs, axis=0, keepdims=True))
+        # print(np.max(self.eig_vecs))
+        self.eig_vecs /= 200 #np.linalg.norm(self.eig_vecs, axis=0, keepdims=True)
+        # np.save("norm_v2.npy", np.max(np.abs(self.eig_vecs), axis=0, keepdims=True))
+        # np.save("norm_v2.npy", np.max(np.abs(self.eig_vecs), axis=0, keepdims=True))
+        # print(np.linalg.norm(self.eig_vecs, axis=0))
+        # print(np.max(self.eig_vecs, axis=0))
+        # print(np.min(self.eig_vecs, axis=0))
+        # print(self.eig_vals)
+
+        # self.eig_vecs /= np.load('norm_v1.npy')
+
+        # self.eig_vecs = self.eig_vecs / np.max(np.abs(self.eig_vecs), axis=0, keepdims=True)
+        # self.eig_vecs = self.eig_vecs * np.load('norm_v2.npy')
+
+
+        # nn = np.load('norm2.npy')
+        # self.eig_vecs /= nn[:,:50]
+
+
+        # ==== VISUALIZATION CODE ==========
+        if False:
+            num_mesh_to_viz = 6
+            meshes = []
+            for i in range(num_mesh_to_viz):
+                meshes.append(trimesh.Trimesh(self.vertices + np.array([i*1,0,0]), self.faces, process=False))
+
+            # mesh = meshes[0].union( meshes[1])
+            # mesh = mesh.union( meshes[2])
+            # mesh = mesh.union( meshes[3])
+            meshes = [trimesh.util.concatenate(meshes)]
+
+            from vedo import trimesh2vedo, show, screenshot, Plotter
+
+            vp = Plotter(axes=0, offscreen=True)
+
+            vmeshes = trimesh2vedo(meshes)
+            cmaps = ('jet', 'PuOr', 'viridis')
+            scals =  self.eig_vecs[:,:num_mesh_to_viz].transpose((1,0)).reshape(-1)
+            vmeshes[0].cmap(cmaps[0], scals).lighting('plastic')
+
+            # add a 2D scalar bar to a mesh
+            vmeshes[0].addScalarBar(title=f"scalarbar #{0}", c='k') 
+
+            vp.show(vmeshes, axes=1)
+            screenshot(f"test_{time()}.png")
+            import sys
+            sys.exit(0)
+        #  ================
+
+
+
+
+        # range between biggest and smallest eigenvalue : 
+        # 6 0.09622419080119388
+        # 6_bis 0.09651935545457718
+        delta = (np.log(self.eig_vals[-1]) - np.log(self.eig_vals[1])) / self.timestamps
+        sigma = 7 * delta
+        e_min = np.log(self.eig_vals[1]) + 2 * delta
+        e_max = np.log(self.eig_vals[-1]) - 2 * delta
+        es = np.linspace(e_min, e_max, self.timestamps)  # T
+        self.delta = delta
+
+        
+        coef = np.expand_dims(es, 0) - np.expand_dims(np.log(self.eig_vals[1:]), 1)  # (K-1)xT
+        coef = np.exp(-np.square(coef) / (2 * sigma * sigma))  # (K-1)xT #element wise square
+        sum_coef = coef.sum(0)  # T
+        K = np.matmul(np.square(self.eig_vecs[:, 1:]), coef)  # VxT. Scaling of the eigen vectors by coef. Coef depends only on the eigen values. Triangulation agnostic.
+        self.wks = K / np.expand_dims(sum_coef, 0)  # VxT Scaling of the eigen vectors by sum_coef. Coef depends only on the eigen values. Triangulation agnostic.
+        # print(np.linalg.norm(self.wks, axis=0))
+        # print(np.linalg.norm(self.wks, axis=1))
+        

NeuralJacobianFields/PoissonSystem.py

@@ -0,0 +1,796 @@
+import numpy
+import igl
+import numpy as np
+import torch
+import time
+
+from scipy.sparse import diags,coo_matrix
+from scipy.sparse import csc_matrix as sp_csc
+
+
+USE_TORCH_SPARSE = True ## This uses TORCH_SPARSE instead of TORCH.SPARSE
+
+# This four are mutually exclusive
+USE_CUPY = False  ## This uses CUPY LU decomposition on GPU
+USE_CHOLESPY_GPU = True  ## This uses cholesky decomposition on GPU
+USE_CHOLESPY_CPU = False  ## This uses cholesky decomposition on CPU
+USE_SCIPY = False ## This uses CUPY LU decomposition on CPU
+
+
+
+# If USE_SCIPY = True, wether or not to use enhanced backend
+USE_SCIKITS_UMFPACK = False  ## This uses UMFPACK backend for scipy instead of naive scipy.
+
+if USE_CHOLESPY_GPU or USE_CHOLESPY_CPU:
+    from cholespy import CholeskySolverD, MatrixType
+
+if USE_CUPY and torch.cuda.is_available():
+    from cupyx.scipy.sparse.linalg import spsolve_triangular
+    from cupyx.scipy.sparse import csr_matrix
+    import cupy
+    from torch.utils.dlpack import to_dlpack, from_dlpack
+
+from scipy.sparse.linalg import splu as scipy_splu
+from scipy.sparse.linalg import spsolve_triangular, spsolve
+if USE_SCIPY:
+    if  USE_SCIKITS_UMFPACK:
+        # This is a bit slower in practice
+        # https://stackoverflow.com/questions/64401503/is-there-a-way-to-further-improve-sparse-solution-times-using-python
+        from scikits.umfpack import splu as scipy_splu
+    else:
+        import scipy.sparse.linalg as lg
+        lg.use_solver(useUmfpack=False)
+        # Slight performance gain with True
+        # conda install -c conda-forge scikit-umfpack
+        # forward pass goes from 0.038 to 0.036
+        # assumeSortedIndices=True Does not bring any boost
+        from scipy.sparse.linalg import splu as scipy_splu
+        from scipy.sparse.linalg import spsolve_triangular, spsolve
+
+
+if USE_TORCH_SPARSE:
+    import torch_sparse
+
+
+USE_UGLY_PATCH_FOR_CUPY_ERROR = False
+
+
+class SparseMat:
+    '''
+    Sparse matrix object represented in the COO format
+    Refacto : consider killing this object, byproduct of torch_sparse instead of torch.sparse (new feature)
+    '''
+
+    @staticmethod
+    def from_M(M,ttype):
+        return SparseMat(M[0],M[1],M[2],M[3],ttype)
+
+    @staticmethod
+    def from_coo(coo,ttype):
+        inds = numpy.vstack((coo.row,coo.col))
+        return SparseMat(inds,coo.data,coo.shape[0],coo.shape[1],ttype)
+
+    def __init__(self,inds,vals,n,m,ttype):
+        self.n = n
+        self.m = m
+        self.vals = vals
+        self.inds = inds
+        assert(inds.shape[0] == 2)
+        assert(inds.shape[1] == vals.shape[0])
+        assert(np.max(inds[0,:]) <= n)
+        assert(np.max(inds[1,:] <= m))
+        #TODO figure out how to extract the I,J,V,m,n from this, then load a COO mat directly from npz
+        #self.coo_mat = coo_matrix((cupy.array(self.vals), (cupy.array(self.inds[0,:]), cupy.array(self.inds[1,:]))))
+        self.vals = torch.from_numpy(self.vals).type(ttype).contiguous()
+        self.inds = torch.from_numpy(self.inds).type(torch.int64).contiguous()
+
+    def to_coo(self):
+        return coo_matrix((self.vals, (self.inds[0,:], self.inds[1,:])), shape = (self.n, self.m))
+
+    def to_csc(self):
+        return sp_csc((self.vals, (self.inds[0,:], self.inds[1,:])), shape = (self.n, self.m))
+
+    def to_cholesky(self):
+        return CholeskySolverD(self.n, self.inds[0,:], self.inds[1,:], self.vals, MatrixType.COO)
+
+    def to(self,device):
+        self.vals = self.vals.to(device)
+        self.inds = self.inds.to(device)
+        return self
+
+    def pin_memory(self):
+        return
+        # self.vals.pin_memory()
+        # self.inds.pin_memory()
+
+    def multiply_with_dense(self,dense):
+        if USE_TORCH_SPARSE:
+            res = torch_sparse.spmm(self.inds,self.vals, self.n, self.m, dense)
+            # 1000 for loop on the above line takes 0.13 sec. Fast but annoying to have this dependency
+        else:
+            # Somehow this is not implemented for now?
+            # res = torch.smm(torch.sparse_coo_tensor(self.inds,self.vals) , (dense.float())).to_dense().to(dense.device)
+            # 1000 for loop on the above line takes 10 sec on the CPU. It is not implemented on gpu yet Slower but no dependency
+            if self.vals.device.type == 'cpu':
+                tensor_zero_hack  = torch.FloatTensor([0]).double() # This line was somehow responsible for a nasty NAN bug
+            else:
+                tensor_zero_hack  =  torch.cuda.FloatTensor([0]).to(dense.get_device()).double()
+            # beware with addmm, it is experimental and gave me a NaN bug!
+            res = torch.sparse.addmm(tensor_zero_hack, torch.sparse_coo_tensor(self.inds.double(),self.vals.double()) , (dense.double())).type_as(self.vals)
+            # 1000 for loop on the above line takes 0.77 sec. Slower but no dependency
+        return res.contiguous()
+
+
+
+class PoissonSystemMatrices:
+    '''
+    Holds the matrices needed to perform gradient and poisson computations
+    Logic : this class is supposed is supposed to hold everything needed to compute Poisson Solver
+    Refacto : merge with Poisson Solver
+    Only accept SparseMat representation
+    '''
+    def __init__(self, V, F,grad, rhs, w, ttype, is_sparse = True, lap = None, cpuonly=False):
+        self.dim = 3
+        self.is_sparse = is_sparse
+        self.w = w
+        self.rhs = rhs
+        self.igl_grad = grad
+        self.ttype = ttype
+        self.__splu_L = None
+        self.__splu_U = None
+        self.__splu_perm_c = None
+        self.__splu_perm_r = None
+        self.lap = lap
+        self.__V = V
+        self.__F = F
+        self.cpuonly = cpuonly
+        self.cpu_splu = None
+    
+
+    def create_poisson_solver(self):
+        return PoissonSolver(self.igl_grad,self.w,self.rhs, None, self.lap)
+
+    def create_poisson_solver_from_splu_old(self, lap_L, lap_U, lap_perm_c, lap_perm_r):
+        w = torch.from_numpy(self.w).type(self.ttype)
+        lap = None
+        my_splu = None
+        if not self.cpuonly:
+            if USE_CUPY:
+                my_splu = MyCuSPLU(lap_L, lap_U, lap_perm_c, lap_perm_r)
+            else:
+                if self.lap is not None:
+                    lap = self.lap
+                    # my_splu = scipy_splu(self.lap)
+                    # my_splu = MyCuSPLU_CPU(lap_L, lap_U, lap_perm_c, lap_perm_r)
+                else: 
+                    my_splu = MyCuSPLU_CPU(lap_L, lap_U, lap_perm_c, lap_perm_r)
+                # st = time.time()
+                # my_splu = scipy_splu(lap_L@lap_U)
+                # print(f"time for LU: {time.time() - st}" )
+
+        else:
+            if self.lap is not None:
+                my_splu = scipy_splu(self.lap)
+            else:
+                0/0
+                # my_splu = splu(lap_L) 
+
+        return PoissonSolver(self.igl_grad,w,self.rhs,my_splu, lap)
+
+    def compute_poisson_solver_from_laplacian(self, compute_splu=True):
+        self.compute_laplacian()
+        if compute_splu:
+            self.compute_splu()
+        return self.create_poisson_solver_from_splu(self.__splu_L,self.__splu_U,self.__splu_perm_c,self.__splu_perm_r)
+
+    def compute_laplacian(self):
+        if self.lap is None:
+            self.lap = igl.cotmatrix(self.__V,self.__F)
+            self.lap = self.lap[1:, 1:]
+            self.lap = SparseMat.from_coo(self.lap.tocoo(), torch.float64)
+        
+        if isinstance(self.lap,PoissonSystemMatrices) and self.lap.vals.shape[0] == self.__V.shape[0]:
+            assert(False), "this should not happen, the fix is to remove a column and row of the laplacian"
+            self.lap = self.lap[1:, 1:]
+
+        return self.lap
+
+    def compute_splu(self):
+        print("i am computing splu")
+        if self.cpu_splu is None:
+            # st = time.time()
+            s = scipy_splu(self.lap)
+            # print(f"time to compute LU {time.time() - st}")
+            # We are storing these attributes just in case we need to create a PoissonSolver on the GPU, they are useless for CPU case.
+            self.cpu_splu = s
+            self.__splu_L = s.L
+            self.__splu_U = s.U
+            self.__splu_perm_c = s.perm_c
+            self.__splu_perm_r = s.perm_r
+        return self.__splu_L,self.__splu_U,self.__splu_perm_c,self.__splu_perm_r
+
+    def get_new_grad(self):
+        grad = self.igl_grad.to_coo()
+        self.igl_grad = SparseMat.from_M(_convert_sparse_igl_grad_to_our_convention(grad.tocsc()),torch.float64)
+        return self.igl_grad
+
+def _convert_sparse_igl_grad_to_our_convention(input):
+    '''
+    The grad operator computed from igl.grad() results in a matrix of shape (3*#tri x #verts).
+    It is packed such that all the x-coordinates are placed first, followed by y and z. As shown below
+
+    ----------           ----------
+    | x1 ...             | x1 ...
+    | x2 ...             | y1 ...
+    | x3 ...             | z1 ...
+    | .                  | .
+    | .                  | .
+    | y1 ...             | x2 ...
+    | y2 ...      ---->  | y2 ...
+    | y3 ...             | z2 ...
+    | .                  | .
+    | .                  | .
+    | z1 ...             | x3 ...
+    | z2 ...             | y3 ...
+    | z3 ...             | z3 ...
+    | .                  | .
+    | .                  | .
+    ----------           ----------
+
+    Note that this functionality cannot be computed trivially if because igl.grad() is a sparse tensor and as such
+    slicing is not well defined for sparse matrices. the following code performs the above conversion and returns a
+    torch.sparse tensor.
+    Set check to True to verify the results by converting the matrices to dense and comparing it.
+    '''
+    assert type(input) == sp_csc, 'Input should be a scipy csc sparse matrix'
+    T = input.tocoo()
+
+    r_c_data = np.hstack((T.row[..., np.newaxis], T.col[..., np.newaxis],
+                          T.data[..., np.newaxis]))  # horizontally stack row, col and data arrays
+    r_c_data = r_c_data[r_c_data[:, 0].argsort()]  # sort along the row column
+
+    # Separate out x, y and z blocks
+    '''
+    Note that for the grad operator there are exactly 3 non zero elements in a row
+    '''
+    L = T.shape[0]
+    Tx = r_c_data[:L, :]
+    Ty = r_c_data[L:2 * L, :]
+    Tz = r_c_data[2 * L:3 * L, :]
+
+    # align the y,z rows with x so that they too start from 0
+    Ty[:, 0] -= Ty[0, 0]
+    Tz[:, 0] -= Tz[0, 0]
+
+    # 'strech' the x,y,z rows so that they can be interleaved.
+    Tx[:, 0] *= 3
+    Ty[:, 0] *= 3
+    Tz[:, 0] *= 3
+
+    # interleave the y,z into x
+    Ty[:, 0] += 1
+    Tz[:, 0] += 2
+
+    Tc = np.zeros((input.shape[0] * 3, 3))
+    Tc[::3] = Tx
+    Tc[1::3] = Ty
+    Tc[2::3] = Tz
+
+    indices = Tc[:, :-1].astype(int)
+    data = Tc[:, -1]
+
+    return (indices.T, data, input.shape[0], input.shape[1])
+
+
+class PoissonSolver:
+    '''
+    an object to compute gradients and solve poisson
+    '''
+
+    def __init__(self,grad,W,rhs,my_splu, lap=None):
+        self.W = torch.from_numpy(W).double()
+        self.grad = grad
+        self.rhs = rhs
+        self.my_splu = my_splu
+        self.lap = lap
+        self.sparse_grad = grad
+        self.sparse_rhs = rhs
+
+    def to(self,device):
+        self.W = self.W.to(device)
+        self.sparse_grad = self.sparse_grad.to(device)
+        self.sparse_rhs = self.sparse_rhs.to(device)
+        if USE_CUPY or USE_CHOLESPY_GPU:
+            self.lap = self.lap.to(device)
+        return self
+
+    def jacobians_from_vertices(self,V):
+        res = _multiply_sparse_2d_by_dense_3d(self.sparse_grad, V).type_as(V)
+        res = res.unsqueeze(2)
+        return res.view(V.shape[0], -1, 3,3).transpose(2,3)
+
+    def restrict_jacobians(self,D):
+        assert isinstance(D, torch.Tensor) and len(D.shape) in [3, 4]
+        assert D.shape[-1] == 3 and D.shape[-2] == 3
+        assert isinstance(self.W, torch.Tensor) and len(self.W.shape) == 3
+        assert self.W.shape[-1] == 2 and self.W.shape[-2] == 3
+
+        if len(D.shape) == 4:
+            DW = torch.einsum("abcd,bde->abce", (D, self.W.type_as(D)))
+        else:
+            DW = torch.einsum("abcd,bde->abce", (D.unsqueeze(0), self.W)).squeeze(0)
+
+        if len(DW.shape)>4:
+            DW = DW.squeeze(0)
+        return DW
+
+    def restricted_jacobians_from_vertices(self,V):
+        return self.restrict_jacobians(self.jacobians_from_vertices(V))
+
+    def solve_poisson(self,jacobians): 
+        # st = time.time()
+        assert(len(jacobians.shape) == 4)
+        assert(jacobians.shape[2] == 3 and jacobians.shape[3] == 3)
+       
+        # torch.cuda.synchronize()
+        # st = time.time()
+
+        if self.my_splu is None:
+            if isinstance(self.lap,SparseMat):
+                # self.my_splu =  scipy_splu(self.lap.to('cpu').to_coo())
+                if USE_CHOLESPY_CPU or USE_CHOLESPY_GPU:
+                    self.my_splu = self.lap.to_cholesky()
+                else:
+                    self.my_splu = scipy_splu(self.lap.to('cpu').to_coo())
+            else:
+                self.my_splu = scipy_splu(self.lap)
+
+        # print(f"computing poisson! {self.lap.vals.get_device()}")
+        # print(f"computing poisson! {self.lap.inds.get_device()}")
+        # print(f"computing poisson! {jacobians.get_device()}")
+        # print(f"computing poisson! {self.sparse_rhs.vals.get_device()}")
+        # torch.cuda.synchronize()
+        # print(f"SOLVER decomposition {time.time() - st}")
+
+        sol = _predicted_jacobians_to_vertices_via_poisson_solve(self.my_splu, self.sparse_rhs, jacobians.transpose(2, 3).reshape(jacobians.shape[0], -1, 3, 1).squeeze(3).contiguous())
+        # torch.cuda.synchronize()
+        # print(f"POISSON LU + SOLVE FORWARD{time.time() - st}")
+        c = torch.mean(sol, axis=1).unsqueeze(1)  ## Beware the predicted mesh is centered here.
+        # print(f"time for poisson: {time.time() - st}" )
+        return sol - c
+
+    def pin_memory(self):
+        return
+        # self.W.pin_memory()
+        # self.sparse_grad.pin_memory()
+        # self.sparse_rhs.pin_memory()
+
+
+def poisson_system_matrices_from_mesh( V,F, dim=3,ttype = torch.float64, is_sparse=True,cpuonly=False):
+    '''
+    compute poisson matricees for a given mesh
+    :param V vertices
+    :param F faces
+    :param dim: for now always 3 :)
+    :param ttype the type of tensor (e.g., float,double)
+    :param is_sparse: for now always true
+    :return: a PoissonMatricese object holding the computed matrices
+    '''
+
+    assert type(dim) == int and dim in [2,3], f'Only two and three dimensional meshes are supported'
+    assert type(is_sparse) == bool
+    vertices = V
+    faces = F
+    dim = 3
+    is_sparse = is_sparse
+
+    grad = igl.grad(vertices, faces)
+    # grad = np.abs(grad)
+    # temp_grad = grad.multiply(csr_matrix(1 / np.sqrt(grad.multiply(grad).sum(1))))
+    # gradients_normalized = grad / np.linalg.norm(grad, axis=1)[:, np.newaxis]
+
+    mass = _get_mass_matrix(vertices,faces,is_sparse)
+    ## TODO 2D Case ##
+    if dim == 2:
+        grad = grad[:-grad.shape[0]//3,:]
+        mass = mass[:-mass.shape[0]//3,:-mass.shape[0]//3]
+
+    laplace = grad.T@mass@grad
+    laplace = laplace[1:, 1:]
+
+    rhs = grad.T@mass
+    b1,b2,_ = igl.local_basis(V,F)
+    w = np.stack((b1,b2),axis=-1)
+    # print(time.time() - s)
+
+    rhs = rhs[1:,:]
+
+    if is_sparse:
+        laplace = laplace.tocoo()
+        rhs = rhs.tocoo()
+        grad = grad.tocsc()
+    else:
+        laplace = laplace.toarray()
+        rhs = rhs.toarray()
+        grad = grad.toarray()
+
+
+    grad = SparseMat.from_M(_convert_sparse_igl_grad_to_our_convention(grad), torch.float64)
+    poissonbuilder =  PoissonSystemMatrices(V=V,F=F,grad=grad, 
+              rhs=SparseMat.from_coo(rhs, torch.float64), w=w,
+              ttype=ttype,is_sparse=is_sparse, 
+              lap=SparseMat.from_coo(laplace, torch.float64), 
+              cpuonly=cpuonly)
+    # poissonbuilder.get_new_grad()
+    return poissonbuilder
+
+def _get_mass_matrix(vertices,faces,is_sparse):
+
+    d_area = igl.doublearea(vertices,faces)
+    d_area = np.hstack((d_area, d_area, d_area))
+    if is_sparse:
+        return sp_csc(diags(d_area))
+    return diags(d_area)
+
+
+
+
+class SPLUSolveLayer(torch.autograd.Function):
+    '''
+    Implements the SPLU solve as a differentiable layer, with a forward and backward function
+    '''
+
+    @staticmethod
+    def forward(ctx, solver, b):
+        '''
+        override forward function
+        :param ctx: context object (to keep the lu object for the backward pass)
+        :param lu: splu object
+        :param b: right hand side, could be a vector or matrix
+        :return: the vector or matrix x which holds lu.solve(b) = x
+        '''
+        assert isinstance(b, torch.Tensor)
+        assert b.shape[-1] >= 1 and b.shape[-1] <= 3, f'got shape {b.shape} expected last dim to be in range 1-3'
+        b = b.contiguous()
+        ctx.solver = solver
+
+        # st = time.time()
+        vertices = SPLUSolveLayer.solve(solver, b).type_as(b)
+        # print(f"FORWARD SOLVE {time.time() - st}")
+
+        assert not torch.isnan(vertices).any(), "Nan in the forward pass of the POISSON SOLVE"
+        return vertices
+
+    def backward(ctx, grad_output):
+        '''
+        overrides backward function
+        :param grad_output: the gradient to be back-propped
+        :return: the outgoing gradient to be back-propped
+        '''
+
+        assert isinstance(grad_output, torch.Tensor)
+        assert grad_output.shape[-1] >= 1 and grad_output.shape[
+            -1] <= 3, f'got shape {grad_output.shape} expected last dim to be in range 1-3'
+        # when backpropping, if a layer is linear with matrix M, x ---> Mx, then the backprop of gradient g is M^Tg
+        # in our case M = A^{-1}, so the backprop is to solve x = A^-T g.
+        # Because A is symmetric we simply solve A^{-1}g without transposing, but this will break if A is not symmetric.
+        # st = time.time()
+        grad_output = grad_output.contiguous()
+        grad = SPLUSolveLayer.solve(ctx.solver,
+                                          grad_output)
+        # print(f"BACKWARD SOLVE {time.time() - st}")
+        # At this point we perform a NAN check because the backsolve sometimes returns NaNs.
+        assert not torch.isnan(grad).any(),  "Nan in the backward pass of the POISSON SOLVE"
+
+        if USE_CUPY:
+            mempool = cupy.get_default_memory_pool()
+            pinned_mempool = cupy.get_default_pinned_memory_pool()
+            mempool.free_all_blocks()
+            pinned_mempool.free_all_blocks()
+            del ctx.lu
+        
+        return None, grad
+
+    @staticmethod
+    def solve(solver, b):
+        '''
+        solve the linear system defined by an SPLU object for a given right hand side. if the RHS is a matrix, solution will also be a matrix.
+        :param solver: the splu object (LU decomposition) or cholesky object
+        :param b: the right hand side to solve for
+        :return: solution x which satisfies Ax = b where A is the poisson system lu describes
+        '''
+
+        if  USE_CUPY:
+            b_cupy = cupy.fromDlpack(to_dlpack(b))
+            with cupy.cuda.Device(solver.device()):
+                # this will hold the solution
+                sol = cupy.ndarray(b_cupy.shape)
+                for i in range(b_cupy.shape[2]):  # b may have multiple columns, solve for each one
+                    b2d = b_cupy[..., i]  # cupy.expand_dims(b_cpu[...,i],2)
+                    s = solver.solve(b2d.T).T
+                    sol[:, :, i] = s
+            # # # convert back to torch
+            res = from_dlpack(sol.toDlpack())
+            # np.save("res_gpu.npy", res.cpu().numpy())
+            # res = torch.zeros((1, 6889, 3), device=b.device)+  torch.mean(b)
+
+            return res.type_as(b.type())
+
+        elif USE_SCIPY:
+            #only CPU
+            # st = time.time()
+            assert(b.shape[0]==1), "Need to code parrallel implem on the first dim"
+            sol = solver.solve(b[0].double().cpu().numpy())
+            res = torch.from_numpy(sol).to(b.device).reshape(b.shape)
+            # print(time.time() - st)
+            return res.type_as(b).contiguous()
+
+            # Legacy code, I don't understand what is the reason for having a for loop
+            # sol = np.ndarray(b.shape)
+            # for i in range(b.shape[2]):  # b may have multiple columns, solve for each one
+            #     b2d = b[..., i]  # cupy.expand_dims(b_cpu[...,i],2)
+            #     s = lu.solve(b2d.double().cpu().float().numpy().T).T
+            #     sol[:, :, i] = s
+            # res = torch.from_numpy(sol).to(b.device)
+            # # np.save("res_cpu.npy", sol)
+            # print(f"time {time.time() - st}" )
+        elif USE_CHOLESPY_GPU:
+            # torch.cuda.synchronize()
+            # # st = time.time()
+            # assert(b.shape[0]==1), "Need to code parrallel implem on the first dim"
+            # b = b.squeeze().double()
+            # x = torch.zeros_like(b)
+            # solver.solve(b, x)
+            # # torch.cuda.synchronize()
+            # # print(f"time cholescky GPU {time.time() - st}" )
+            # return x.contiguous().unsqueeze(0)
+            # st = time.time()
+            # print(b.get_device(), b.shape)
+            b = b.double().contiguous()
+            c = b.permute(1,2,0).contiguous()
+            c = c.view(c.shape[0], -1)
+            x = torch.zeros_like(c)
+            solver.solve(c, x)
+            x = x.view(b.shape[1], b.shape[2], b.shape[0])
+            x = x.permute(2,0,1).contiguous()
+            # torch.cuda.synchronize()
+            # print(f"time cholescky GPU {time.time() - st}" )
+            return x.contiguous()
+
+        elif USE_CHOLESPY_CPU:
+            # st = time.time()
+            assert(b.shape[0]==1), "Need to code parrallel implem on the first dim"
+            b = b.squeeze()
+            b_cpu = b.cpu()
+            x = torch.zeros_like(b_cpu)
+            solver.solve(b_cpu, x)
+            # print(f"time cholescky CPU {time.time() - st}" )
+            return x.contiguous().to(b.device).unsqueeze(0)
+
+
+        return res.type_as(b)
+
+def _predicted_jacobians_to_vertices_via_poisson_solve(Lap, rhs, jacobians):
+    '''
+    convert the predictions to the correct convention and feed it to the poisson solve
+    '''
+
+    def _batch_rearrange_input(input):
+        assert isinstance(input, torch.Tensor) and len(input.shape) in [2, 3]
+        P = torch.zeros(input.shape).type_as(input)
+        if len(input.shape) == 3:
+            # Batched input
+            k = input.shape[1] // 3
+            P[:, :k, :] = input[:, ::3]
+            P[:, k:2 * k, :] = input[:, 1::3]
+            P[:, 2 * k:, :] = input[:, 2::3]
+
+        else:
+            k = input.shape[0] // 3
+            P[:k, :] = input[::3]
+            P[k:2 * k, :] = input[1::3]
+            P[2 * k:, :] = input[2::3]
+
+        return P
+
+    def _list_rearrange_input(input):
+        assert isinstance(input, list) and all([isinstance(x, torch.Tensor) and len(x.shape) in [2, 3] for x in input])
+        P = []
+        for p in input:
+            P.append(_batch_rearrange_input(p))
+        return P
+
+    if isinstance(jacobians, list):
+        P = _list_rearrange_input(jacobians)
+    else:
+        P = _batch_rearrange_input(jacobians)
+
+    # return solve_poisson(Lap, rhs, P)
+    assert isinstance(P, torch.Tensor) and len(P.shape) in [2, 3]
+    assert len(P.shape) == 3
+
+    # torch.cuda.synchronize()
+    # st = time.time()
+    P = P.double()
+    input_to_solve = _multiply_sparse_2d_by_dense_3d(rhs, P)
+
+
+    out = SPLUSolveLayer.apply(Lap, input_to_solve)
+
+    out = torch.cat([torch.zeros(out.shape[0], 1, out.shape[2]).type_as(out), out], dim=1)  ## Why?? Because!
+    out = out - torch.mean(out, axis=1, keepdim=True)
+
+    return out.type_as(jacobians)
+
+
+
+def _multiply_sparse_2d_by_dense_3d(mat, B):
+    ret = []
+    for i in range(B.shape[0]):
+        C = mat.multiply_with_dense(B[i, ...])
+        ret.append(C)
+    ret = torch.stack(tuple(ret))
+    return ret
+
+
+
+
+
+
+
+class MyCuSPLU:
+    '''
+    implmentation of SPLU on the gpu via CuPy
+    '''
+    def __init__(self, L, U, perm_c=None, perm_r=None):
+        # with cupy.cuda.Device(device):
+        self.__orgL = L
+        self.__orgU = U
+        # self.L = csr_matrix(L)
+        # self.U = csr_matrix(U)
+        self.L = None
+        self.U = None
+        self.perm_c = perm_c
+        self.perm_r = perm_r
+        # self.splu = cu_splu(csr_matrix(lap))
+        # self.L = self.splu.L
+        # self.U = self.splu.U
+        # self.perm_c = self.splu.perm_c
+        # self.perm_r = self.splu.perm_r
+        self.__device = None
+
+    def to(self, device):
+        # assumes to receive a pytorch device object that has a "index" field
+        # print(device)
+        # if(self.__device is None):
+        #     raise Exception()
+        self.__device = device.index
+        with cupy.cuda.Device(self.__device):
+            # self.__orgL = cupy.asarray(self.__orgL)
+            # self.__orgU = cupy.asarray(self.__orgU)
+            self.L = csr_matrix(self.__orgL)
+            self.U = csr_matrix(self.__orgU)
+        return self
+
+    def device(self):
+        return self.__device
+
+    def solve(self, b):
+        """ an attempt to use SuperLU data to efficiently solve
+            Ax = Pr.T L U Pc.T x = b
+             - note that L from SuperLU is in CSC format solving for c
+               results in an efficiency warning
+            Pr . A . Pc = L . U
+            Lc = b      - forward solve for c
+             c = Ux     - then back solve for x
+        """
+
+        assert self.__device is not None, "need to explicitly call to() before solving"
+        if USE_UGLY_PATCH_FOR_CUPY_ERROR:
+            with cupy.cuda.Device(0):
+                b[:1, :1].copy()[:, :1]
+
+        with cupy.cuda.Device(self.__device):
+            b = cupy.array(b)
+            if self.perm_r is not None:
+                b_old = b.copy()
+                b[self.perm_r] = b_old
+
+        assert b.device.id == self.__device, "got device" + str(b.device.id) + "instead of" + str(self.__device)
+        # st = time.time()
+        try:  # unit_diagonal is a new kw
+            c = spsolve_triangular(self.L, b, lower=True, unit_diagonal=True, overwrite_b=True)
+        except TypeError:
+            c = spsolve_triangular(self.L, b, lower=True, overwrite_b=True)
+        px = spsolve_triangular(self.U, c, lower=False, overwrite_b=True)
+        # print(f"time for spsolve_triangular GPU: {time.time() - st}" )
+
+        if self.perm_c is None:
+            return px
+        px = px[self.perm_c]
+
+        # print(f'used: {mempool.used_bytes()}')
+        # print(f'total: {mempool.total_bytes()}')
+        return px
+
+
+class MyCuSPLU_CPU:
+    '''
+    implmentation of SPLU on the gpu via CuPy
+    '''
+    def __init__(self, L, U, perm_c=None, perm_r=None):
+        # with cupy.cuda.Device(device):
+        self.__orgL = L
+        self.__orgU = U
+        # self.L = csr_matrix(L)
+        # self.U = csr_matrix(U)
+        self.L = L
+        self.U = U
+        # self.L = L.tocsr()
+        # self.U = U.tocsr()
+        self.perm_c = perm_c
+        self.perm_r = perm_r
+        # self.splu = cu_splu(csr_matrix(lap))
+        # self.L = self.splu.L
+        # self.U = self.splu.U
+        # self.perm_c = self.splu.perm_c
+        # self.perm_r = self.splu.perm_r
+        self.__device = 'cpu'
+
+    def to(self, device):
+        # assumes to receive a pytorch device object that has a "index" field
+        # print(device)
+        # if(self.__device is None):
+        #     raise Exception()
+        # self.__device = device.index
+        # with cupy.cuda.Device(self.__device):
+        #     # self.__orgL = cupy.asarray(self.__orgL)
+        #     # self.__orgU = cupy.asarray(self.__orgU)
+        #     self.L = csr_matrix(self.__orgL)
+        #     self.U = csr_matrix(self.__orgU)
+        return self
+
+    def device(self):
+        return self.__device
+
+    def solve(self, b):
+        """ an attempt to use SuperLU data to efficiently solve
+            Ax = Pr.T L U Pc.T x = b
+             - note that L from SuperLU is in CSC format solving for c
+               results in an efficiency warning
+            Pr . A . Pc = L . U
+            Lc = b      - forward solve for c
+             c = Ux     - then back solve for x
+        """
+
+
+        # Could be done on GPU
+        if self.perm_r is not None:
+            b_old = b.copy()
+            b[self.perm_r] = b_old
+        # ,  permc_spec="NATURAL"
+        # ,  permc_spec="NATURAL"
+        # ,  permc_spec="NATURAL"
+        st = time.time()
+        # try:  # unit_diagonal is a new kw
+        #     c = spsolve_triangular(self.L, b, lower=True, unit_diagonal=True, overwrite_b=True)
+        # except TypeError:
+        #     c = spsolve_triangular(self.L, b, lower=True, overwrite_b=True)
+        # px = spsolve_triangular(self.U, c, lower=False, overwrite_b=True)
+        try:  # unit_diagonal is a new kw
+            c = spsolve(self.L, b, permc_spec="NATURAL")
+        except TypeError:
+            c = spsolve(self.L, b, permc_spec="NATURAL")
+        px = spsolve(self.U, c, permc_spec="NATURAL")
+        # # (self.L *  c) - b / np.norm(b)
+        print(f"time for spsolve_triangular CPU: {time.time() - st}" )
+
+        if self.perm_c is None:
+            return px
+        px = px[self.perm_c]
+
+        # print(f'used: {mempool.used_bytes()}')
+        # print(f'total: {mempool.total_bytes()}')
+        return px
+
+        # return cupy.asnumpy(px)

NeuralJacobianFields/SourceMesh.py

@@ -0,0 +1,152 @@
+import os
+
+import numpy
+import torch
+import igl
+from . import MeshProcessor
+WKS_DIM = MeshProcessor.WKS_DIM
+WKS_FACTOR = 1000
+import numpy as np
+import sys
+import random
+import time
+class SourceMesh:
+    '''
+    datastructure for the source mesh to be mapped
+    '''
+
+    def __init__(self, source_ind, source_dir, extra_source_fields,
+                 random_scale, ttype, use_wks=False, random_centering=False,
+                cpuonly=False):
+        self.__use_wks = use_wks
+        self.source_ind = source_ind
+        self.source_dir = source_dir
+        self.centroids_and_normals = None
+        self.center_source = True
+        self.poisson = None
+        self.splu = None
+        self.__source_global_translation_to_original = 0
+        self.__extra_keys = extra_source_fields
+        self.__loaded_data = {}
+        self.__ttype = ttype
+        self.__random_scale = random_scale
+        self.random_centering = random_centering
+        self.source_mesh_centroid = None
+        self.mesh_processor = None
+        self.cpuonly = cpuonly
+
+    def get_vertices(self):
+        return self.source_vertices
+
+    def get_global_translation_to_original(self):
+        return self.__source_global_translation_to_original
+
+    def vertices_from_jacobians(self, d):
+        return self.poisson.solve_poisson(d)
+        # return self.splu.solve(d)
+    def jacobians_from_vertices(self, v):
+        return self.poisson.jacobians_from_vertices(v)
+
+    def restrict_jacobians(self, J):
+        return self.poisson.restrict_jacobians(J)
+
+    def get_loaded_data(self, key: str):
+
+        return self.__loaded_data.get(key)
+
+    def get_source_triangles(self):
+        # if self.__source_triangles is None:
+        #     self.__source_triangles = np.load(os.path.join(self.source_dir, 'faces.npy'))
+        return self.mesh_processor.get_faces()
+
+    def to(self, device):
+        self.poisson = self.poisson.to(device)
+        self.splu = self.splu.to(device)
+        self.centroids_and_normals = self.centroids_and_normals.to(device)
+        for key in self.__loaded_data.keys():
+            self.__loaded_data[key] = self.__loaded_data[key].to(device)
+        return self
+
+    def __init_from_mesh_data(self):
+        assert self.mesh_processor is not None
+        self.mesh_processor.prepare_differential_operators_for_use(self.__ttype) #call 1
+        self.source_vertices = torch.from_numpy(self.mesh_processor.get_vertices()).type(
+            self.__ttype)
+        if self.__random_scale != 1:
+            print("Diff ops and WKS need to be multiplied accordingly. Not implemented for now")
+            sys.exit()
+        self.source_vertices *= self.__random_scale
+
+        bb = igl.bounding_box(self.source_vertices.numpy())[0]
+        diag = igl.bounding_box_diagonal(self.source_vertices.numpy())
+
+        # self.source_mesh_centroid = torch.mean(self.source_vertices, axis=0)
+        self.source_mesh_centroid =  (bb[0] + bb[-1])/2
+        if self.random_centering:
+            # centering augmentation
+            self.source_mesh_centroid =  self.source_mesh_centroid + [(2*random.random() - 1)*diag*0.2, (2*random.random() - 1)*diag*0.2, (2*random.random() - 1)*diag*0.2]
+        # self.source_mesh_centroid =  (bb[0] + bb[-1])/2 - np.array([-0.00033245, -0.2910367 ,  0.02100835])
+
+        # Load input to NJF MLP
+        # start = time.time()
+        centroids = self.mesh_processor.get_centroids()
+        centroid_points_and_normals = centroids.points_and_normals
+        if self.__use_wks:
+            wks = WKS_FACTOR * centroids.wks
+            centroid_points_and_normals = numpy.hstack((centroid_points_and_normals, wks))
+        self.centroids_and_normals = torch.from_numpy(
+            centroid_points_and_normals).type(self.__ttype)
+        if self.center_source:
+            c = self.source_mesh_centroid
+            self.centroids_and_normals[:, 0:3] -= c
+            self.source_vertices -= c
+            self.__source_global_translation_to_original = c
+        self.poisson = self.mesh_processor.diff_ops.poisson_solver
+        self.splu = self.mesh_processor.diff_ops.MyCuSPLU_solver
+
+
+        # Essentially here we load pointnet data and apply the same preprocessing
+        for key in self.__extra_keys:
+            data = self.mesh_processor.get_data(key)
+            # if data is None:  # not found in mesh data so try loading from disk
+            #     data = np.load(os.path.join(self.source_dir, key + ".npy"))
+            data = torch.from_numpy(data)
+            if key == 'samples':
+                if self.center_source:
+                    data -= self.get_mesh_centroid()
+                scale = self.__random_scale
+                data *= scale
+            data = data.unsqueeze(0).type(self.__ttype)
+            
+            self.__loaded_data[key] = data
+        # print("Ellapsed load source mesh ", time.time() - start)
+
+    def load(self, source_v=None, source_f=None):
+        # mesh_data = SourceMeshData.SourceMeshData.meshprocessor_from_file(self.source_dir)
+        if source_v is not None and source_f is not None:
+            self.mesh_processor = MeshProcessor.MeshProcessor.meshprocessor_from_array(source_v,source_f, self.source_dir, self.__ttype, cpuonly=self.cpuonly, load_wks_samples=self.__use_wks, load_wks_centroids=self.__use_wks)
+        else:
+            if os.path.isdir(self.source_dir):
+                self.mesh_processor = MeshProcessor.MeshProcessor.meshprocessor_from_directory(self.source_dir, self.__ttype, cpuonly=self.cpuonly, load_wks_samples=self.__use_wks, load_wks_centroids=self.__use_wks)
+            else:
+                self.mesh_processor = MeshProcessor.MeshProcessor.meshprocessor_from_file(self.source_dir, self.__ttype, cpuonly=self.cpuonly, load_wks_samples=self.__use_wks, load_wks_centroids=self.__use_wks)
+        self.__init_from_mesh_data()
+
+    def get_point_dim(self):
+        return self.centroids_and_normals.shape[1]
+
+    def get_centroids_and_normals(self):
+        return self.centroids_and_normals
+
+    def get_mesh_centroid(self):
+        return self.source_mesh_centroid
+
+    def pin_memory(self):
+        # self.poisson.pin_memory()
+        # self.centroids_and_normals.pin_memory()
+        # self.source_vertices.pin_memory()
+        # for key in self.__loaded_data.keys():
+        #     self.__loaded_data[key].pin_memory()
+        return self
+
+

app.py

@@ -0,0 +1,230 @@
+import os
+import sys
+import yaml
+import torch
+import random
+import numpy as np
+import gradio as gr
+from pathlib import Path
+import tempfile
+import shutil
+
+# Add the current directory to Python path
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+# Add packages directory to Python path
+packages_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'packages')
+if os.path.exists(packages_dir):
+    sys.path.append(packages_dir)
+
+try:
+    from loop import loop
+except ImportError as e:
+    print(f"Error importing loop: {e}")
+    print("Make sure all dependencies are installed correctly")
+    sys.exit(1)
+
+# Global variables for configuration
+DEFAULT_CONFIG = {
+    'output_path': './outputs',
+    'gpu': 0,
+    'seed': 99,
+    'clip_model': 'ViT-B/32',
+    'consistency_clip_model': 'ViT-B/32',
+    'consistency_vit_stride': 8,
+    'consistency_vit_layer': 11,
+    'mesh': './meshes/longsleeve.obj',
+    'target_mesh': './meshes_target/jacket_sdf_new.obj',
+    'retriangulate': 0,
+    'bsdf': 'diffuse',
+    'lr': 0.0025,
+    'epochs': 1800,
+    'clip_weight': 2.5,
+    'delta_clip_weight': 5,
+    'vgg_weight': 0.0,
+    'face_weight': 0,
+    'regularize_jacobians_weight': 0.15,
+    'consistency_loss_weight': 0,
+    'consistency_elev_filter': 30,
+    'consistency_azim_filter': 20,
+    'batch_size': 24,
+    'train_res': 512,
+    'resize_method': 'cubic',
+    'fov_min': 30.0,
+    'fov_max': 90.0,
+    'dist_min': 2.5,
+    'dist_max': 3.5,
+    'light_power': 5.0,
+    'elev_alpha': 1.0,
+    'elev_beta': 5.0,
+    'elev_max': 60.0,
+    'azim_min': 0.0,
+    'azim_max': 360.0,
+    'aug_loc': 1,
+    'aug_light': 1,
+    'aug_bkg': 0,
+    'adapt_dist': 1,
+    'log_interval': 5,
+    'log_interval_im': 150,
+    'log_elev': 0,
+    'log_fov': 60.0,
+    'log_dist': 3.0,
+    'log_res': 512,
+    'log_light_power': 3.0
+}
+
+def process_garment(text_prompt, base_text_prompt, epochs, learning_rate, clip_weight, delta_clip_weight, progress=gr.Progress()):
+    """
+    Main function to process garment generation
+    """
+    try:
+        # Create a temporary output directory
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Update configuration
+            config = DEFAULT_CONFIG.copy()
+            config.update({
+                'output_path': temp_dir,
+                'text_prompt': text_prompt,
+                'base_text_prompt': base_text_prompt,
+                'epochs': int(epochs),
+                'lr': float(learning_rate),
+                'clip_weight': float(clip_weight),
+                'delta_clip_weight': float(delta_clip_weight),
+                'gpu': 0  # Use first GPU
+            })
+            
+            # Set random seeds
+            random.seed(config['seed'])
+            os.environ['PYTHONHASHSEED'] = str(config['seed'])
+            np.random.seed(config['seed'])
+            torch.manual_seed(config['seed'])
+            torch.cuda.manual_seed(config['seed'])
+            torch.backends.cudnn.deterministic = True
+            
+            progress(0.1, desc="Initializing...")
+            
+            # Run the main processing loop
+            loop(config)
+            
+            progress(0.9, desc="Processing complete, preparing output...")
+            
+            # Look for output files
+            output_files = []
+            for file_path in Path(temp_dir).rglob("*"):
+                if file_path.is_file() and file_path.suffix.lower() in ['.obj', '.png', '.jpg', '.jpeg', '.gif', '.mp4']:
+                    output_files.append(str(file_path))
+            
+            if output_files:
+                return output_files[0] if len(output_files) == 1 else output_files
+            else:
+                return "Processing completed but no output files found."
+                
+    except Exception as e:
+        return f"Error during processing: {str(e)}"
+
+def create_interface():
+    """
+    Create the Gradio interface
+    """
+    with gr.Blocks(title="Garment3DGen - 3D Garment Stylization", theme=gr.themes.Soft()) as interface:
+        gr.Markdown("""
+        # Garment3DGen: 3D Garment Stylization and Texture Generation
+        
+        This tool allows you to stylize 3D garments using text prompts. Upload a 3D mesh and describe the desired style to generate a new 3D garment.
+        
+        ## How to use:
+        1. Enter a text prompt describing the target style (e.g., "leather jacket with studs")
+        2. Enter a base text prompt describing the input mesh (e.g., "simple t-shirt")
+        3. Adjust the parameters as needed
+        4. Click "Generate" to start the process
+        
+        **Note:** Processing may take several minutes depending on the number of epochs.
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### Input Parameters")
+                
+                text_prompt = gr.Textbox(
+                    label="Target Text Prompt",
+                    placeholder="e.g., leather jacket with studs, denim jacket with patches",
+                    value="leather jacket with studs"
+                )
+                
+                base_text_prompt = gr.Textbox(
+                    label="Base Text Prompt",
+                    placeholder="e.g., simple t-shirt, basic long sleeve shirt",
+                    value="simple t-shirt"
+                )
+                
+                epochs = gr.Slider(
+                    minimum=100,
+                    maximum=3000,
+                    value=1800,
+                    step=100,
+                    label="Number of Epochs",
+                    info="More epochs = better quality but longer processing time"
+                )
+                
+                learning_rate = gr.Slider(
+                    minimum=0.0001,
+                    maximum=0.01,
+                    value=0.0025,
+                    step=0.0001,
+                    label="Learning Rate"
+                )
+                
+                clip_weight = gr.Slider(
+                    minimum=0.1,
+                    maximum=10.0,
+                    value=2.5,
+                    step=0.1,
+                    label="CLIP Weight"
+                )
+                
+                delta_clip_weight = gr.Slider(
+                    minimum=0.1,
+                    maximum=20.0,
+                    value=5.0,
+                    step=0.1,
+                    label="Delta CLIP Weight"
+                )
+                
+                generate_btn = gr.Button("Generate 3D Garment", variant="primary")
+            
+            with gr.Column(scale=1):
+                gr.Markdown("### Output")
+                output = gr.File(label="Generated 3D Garment")
+                status = gr.Textbox(label="Status", interactive=False)
+        
+        # Connect the button to the processing function
+        generate_btn.click(
+            fn=process_garment,
+            inputs=[text_prompt, base_text_prompt, epochs, learning_rate, clip_weight, delta_clip_weight],
+            outputs=[output]
+        )
+        
+        gr.Markdown("""
+        ## Tips for better results:
+        - Be specific in your text prompts
+        - Use descriptive terms for materials, colors, and styles
+        - The base text prompt should accurately describe your input mesh
+        - Higher epoch counts generally produce better results but take longer
+        - Experiment with different CLIP weights for different effects
+        
+        ## Technical Details:
+        This tool uses Neural Jacobian Fields and CLIP embeddings to deform and stylize 3D garment meshes.
+        The process involves optimizing the mesh geometry and texture to match the target text description.
+        """)
+    
+    return interface
+
+if __name__ == "__main__":
+    # Create and launch the interface
+    interface = create_interface()
+    interface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        debug=True
+    ) 

example_config.yml

@@ -0,0 +1,56 @@
+output_path: ./outputs
+gpu: 0
+seed: 99
+
+# CLIP-related
+clip_model: ViT-B/32
+consistency_clip_model: ViT-B/32
+consistency_vit_stride: 8
+consistency_vit_layer: 11
+
+# Mesh
+mesh: ./meshes/longsleeve.obj 
+target_mesh: ./meshes_target/jacket_sdf_new.obj
+retriangulate: 0
+
+# Render settings
+bsdf: diffuse
+
+# Hyper-parameters
+lr: 0.0025 # 0.0025
+epochs: 1800
+clip_weight: 2.5 #0.5 #20.0 # 20.0 #2 for garments
+delta_clip_weight: 5 #10.0 # 2 for garments
+vgg_weight: 0.0
+face_weight: 0
+regularize_jacobians_weight: 0.15 # 0.15 #0.5
+consistency_loss_weight: 0 #-0.25 #-0.15 #-0.5  # -0.25 #-0.5 # -0.25 for garments
+consistency_elev_filter: 30
+consistency_azim_filter: 20
+batch_size: 24
+train_res: 512
+resize_method: cubic
+
+# Camera parameters
+fov_min: 30.0
+fov_max: 90.0
+dist_min: 2.5
+dist_max: 3.5
+light_power: 5.0
+elev_alpha: 1.0
+elev_beta: 5.0
+elev_max: 60.0
+azim_min: 0.0
+azim_max: 360.0
+aug_loc: 1
+aug_light: 1
+aug_bkg: 0
+adapt_dist: 1
+
+log_interval: 5
+log_interval_im: 150
+log_elev: 0
+log_fov: 60.0
+log_dist: 3.0
+log_res: 512
+log_light_power: 3.0

get_embeddings.py

@@ -0,0 +1,58 @@
+import torch
+from PIL import Image
+
+
+def get_fashion_text_embeddings(fclip, cfg, device):
+    print(f'Target text prompt is {cfg.text_prompt}')
+    print(f'Base text prompt is {cfg.base_text_prompt}')
+    with torch.no_grad():
+        text_embeds = fclip.encode_text_tensors([cfg.text_prompt]).detach()
+        base_text_embeds = fclip.encode_text_tensors([cfg.base_text_prompt]).detach()
+        target_text_embeds = text_embeds.clone() / text_embeds.norm(dim=1, keepdim=True)
+        delta_text_embeds = text_embeds - base_text_embeds
+        delta_text_embeds = delta_text_embeds / delta_text_embeds.norm(dim=1, keepdim=True)
+    return target_text_embeds.to(device), delta_text_embeds.to(device)
+
+
+def get_fashion_img_embeddings(fclip, cfg, device, normalize=True):
+    print(f'Target image path is {cfg.image_prompt}')
+    print(f'Base image path is {cfg.base_image_prompt}')
+    with torch.no_grad():
+        target_image_embeds = fclip.encode_images([cfg.image_prompt], 1)
+        target_image_embeds = torch.tensor(target_image_embeds, device=device).detach()
+
+        base_image_embeds = fclip.encode_images([cfg.base_image_prompt], 1)
+        base_image_embeds = torch.tensor(base_image_embeds, device=device).detach()
+        delta_img_embeds = target_image_embeds - base_image_embeds
+    if normalize:
+        delta_img_embeds = delta_img_embeds / delta_img_embeds.norm(dim=1, keepdim=True)
+        target_image_embeds = target_image_embeds.clone() / target_image_embeds.norm(dim=1, keepdim=True)
+    return target_image_embeds.to(device), delta_img_embeds.to(device)
+
+
+def get_text_embeddings(clip, model, cfg, device):
+    print(f'Target text prompt is {cfg.text_prompt}')
+    print(f'Base text prompt is {cfg.base_text_prompt}')
+    text_embeds = clip.tokenize(cfg.text_prompt).to(device)
+    base_text_embeds = clip.tokenize(cfg.base_text_prompt).to(device)
+    with torch.no_grad():
+        text_embeds = model.encode_text(text_embeds).detach()
+        target_text_embeds = text_embeds.clone() / text_embeds.norm(dim=1, keepdim=True)
+        delta_text_embeds = text_embeds - model.encode_text(base_text_embeds)
+        delta_text_embeds = delta_text_embeds / delta_text_embeds.norm(dim=1, keepdim=True)
+    return target_text_embeds, delta_text_embeds
+
+
+def get_img_embeddings(model, preprocess, cfg, device):
+    print(f'Target image path is {cfg.image_prompt}')
+    print(f'Base image path is {cfg.base_image_prompt}')
+
+    image = preprocess(Image.open(cfg.image_prompt)).unsqueeze(0).to(device)
+    base_image = preprocess(Image.open(cfg.base_image_prompt)).unsqueeze(0).to(device)
+    with torch.no_grad():
+        target_image_embeds = model.encode_image(image).to(device).detach()
+        base_image_embeds = model.encode_image(base_image).to(device)
+
+        delta_img_embeds = target_image_embeds - base_image_embeds
+        delta_img_embeds = delta_img_embeds / delta_img_embeds.norm(dim=1, keepdim=True)
+    return target_image_embeds, delta_img_embeds

loop.py

@@ -0,0 +1,436 @@
+import clip
+import kornia
+import os
+import sys
+import pathlib
+import torchvision
+import logging
+import yaml
+import nvdiffrast.torch as dr
+from easydict import EasyDict
+
+from NeuralJacobianFields import SourceMesh
+
+from nvdiffmodeling.src import render
+
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+from utilities.video import Video
+from utilities.helpers import cosine_avg, create_scene, l1_avg
+from utilities.camera import CameraBatch, get_camera_params
+from utilities.clip_spatial import CLIPVisualEncoder
+from utilities.resize_right import resize, cubic, linear, lanczos2, lanczos3
+from packages.fashion_clip.fashion_clip.fashion_clip import FashionCLIP
+from utils import *
+from get_embeddings import *
+
+from pytorch3d.structures import Meshes
+from pytorch3d.loss import (
+    chamfer_distance,
+    mesh_edge_loss,
+    mesh_laplacian_smoothing,
+    mesh_normal_consistency,
+)
+from pytorch3d.ops import sample_points_from_meshes
+
+
+def total_triangle_area(vertices):
+    # Calculate the sum of the areas of all triangles in the mesh
+    num_triangles = vertices.shape[0] // 3
+    triangle_vertices = vertices.view(num_triangles, 3, 3)
+
+    # Calculate the cross product for each triangle
+    cross_products = torch.cross(triangle_vertices[:, 1] - triangle_vertices[:, 0],
+                                 triangle_vertices[:, 2] - triangle_vertices[:, 0])
+
+    # Calculate the area of each triangle
+    areas = 0.5 * torch.norm(cross_products, dim=1)
+
+    # Sum the areas of all triangles
+    total_area = torch.sum(areas)
+    return total_area
+
+def triangle_size_regularization(vertices):
+    # Penalize small triangles by minimizing the squared sum of triangle areas
+    return total_triangle_area(vertices)**2
+
+def loop(cfg):
+    clip_flag = True
+    output_path = pathlib.Path(cfg['output_path'])
+    os.makedirs(output_path, exist_ok=True)
+    with open(output_path / 'config.yml', 'w') as f:
+        yaml.dump(cfg, f, default_flow_style=False)
+    cfg = EasyDict(cfg)
+    
+    print(f'Output directory {cfg.output_path} created')
+    os.makedirs(output_path / 'tmp', exist_ok=True)
+
+    device = torch.device(f'cuda:{cfg.gpu}')
+    torch.cuda.set_device(device)
+
+    text_input, image_input, fashion_image, fashion_text, use_target_mesh = False, False, False, False, True
+    CLIP_embeddings = False
+
+    if CLIP_embeddings:
+        print('Loading CLIP Models')
+        model, preprocess = clip.load(cfg.clip_model, device=device)
+    else:
+        fclip = FashionCLIP('fashion-clip')
+
+    fe = CLIPVisualEncoder(cfg.consistency_clip_model, cfg.consistency_vit_stride, device)
+
+    if fashion_text or fashion_image:
+        target_direction_embeds, delta_direction_embeds = get_fashion_img_embeddings(fclip, cfg, device, True)
+    elif text_input:
+        target_direction_embeds, delta_direction_embeds = get_text_embeddings(clip, model, cfg, device)
+    elif image_input:
+        target_direction_embeds, delta_direction_embeds = get_img_embeddings(model, preprocess, cfg, device)
+
+    clip_mean = torch.tensor([0.48154660, 0.45782750, 0.40821073], device=device)
+    clip_std = torch.tensor([0.26862954, 0.26130258, 0.27577711], device=device)
+
+    # output video
+    video = Video(cfg.output_path)
+    # GL Context
+    glctx = dr.RasterizeGLContext()
+
+
+    load_mesh = get_mesh(cfg.mesh, output_path, cfg.retriangulate, cfg.bsdf)
+
+    if use_target_mesh:
+        target_mesh = get_mesh(cfg.target_mesh, output_path, cfg.retriangulate, cfg.bsdf, 'mesh_target.obj')
+        # We construct a Meshes structure for the target mesh
+        trg_mesh_p3d = Meshes(verts=[target_mesh.v_pos], faces=[target_mesh.t_pos_idx])
+
+
+    jacobian_source = SourceMesh.SourceMesh(0, str(output_path / 'tmp' / 'mesh.obj'), {}, 1, ttype=torch.float)
+    if len(list((output_path / 'tmp').glob('*.npz'))) > 0:
+        logging.warn(f'Using existing Jacobian .npz files in {str(output_path)}/tmp/ ! Please check if this is intentional.')
+    jacobian_source.load()
+    jacobian_source.to(device)
+
+    with torch.no_grad():
+        gt_jacobians = jacobian_source.jacobians_from_vertices(load_mesh.v_pos.unsqueeze(0))
+    gt_jacobians.requires_grad_(True)
+
+    optimizer = torch.optim.Adam([gt_jacobians], lr=cfg.lr)
+    cams_data = CameraBatch(
+        cfg.train_res,
+        [cfg.dist_min, cfg.dist_max],
+        [cfg.azim_min, cfg.azim_max],
+        [cfg.elev_alpha, cfg.elev_beta, cfg.elev_max],
+        [cfg.fov_min, cfg.fov_max],
+        cfg.aug_loc,
+        cfg.aug_light,
+        cfg.aug_bkg,
+        cfg.batch_size,
+        rand_solid=True
+    )
+    cams = torch.utils.data.DataLoader(cams_data, cfg.batch_size, num_workers=0, pin_memory=True)
+    best_losses = {'CLIP': np.inf, 'total': np.inf}
+
+    for out_type in ['final', 'best_clip', 'best_total', 'target_final']:
+        os.makedirs(output_path / f'mesh_{out_type}', exist_ok=True)
+    os.makedirs(output_path / 'images', exist_ok=True)
+    logger = SummaryWriter(str(output_path / 'logs'))
+
+    rot_ang = 0.0
+    t_loop = tqdm(range(cfg.epochs), leave=False)
+
+    if cfg.resize_method == 'cubic':
+        resize_method = cubic
+    elif cfg.resize_method == 'linear':
+        resize_method = linear
+    elif cfg.resize_method == 'lanczos2':
+        resize_method = lanczos2
+    elif cfg.resize_method == 'lanczos3':
+        resize_method = lanczos3
+
+    for it in t_loop:
+
+        # updated vertices from jacobians
+        n_vert = jacobian_source.vertices_from_jacobians(gt_jacobians).squeeze()
+
+        # TODO: More texture code required to make it work ...
+        ready_texture = texture.Texture2D(
+            kornia.filters.gaussian_blur2d(
+                load_mesh.material['kd'].data.permute(0, 3, 1, 2),
+                kernel_size=(7, 7),
+                sigma=(3, 3),
+            ).permute(0, 2, 3, 1).contiguous()
+        )
+
+        kd_notex = texture.Texture2D(torch.full_like(ready_texture.data, 0.5))
+
+        ready_specular = texture.Texture2D(
+            kornia.filters.gaussian_blur2d(
+                load_mesh.material['ks'].data.permute(0, 3, 1, 2),
+                kernel_size=(7, 7),
+                sigma=(3, 3),
+            ).permute(0, 2, 3, 1).contiguous()
+        )
+
+        ready_normal = texture.Texture2D(
+            kornia.filters.gaussian_blur2d(
+                load_mesh.material['normal'].data.permute(0, 3, 1, 2),
+                kernel_size=(7, 7),
+                sigma=(3, 3),
+            ).permute(0, 2, 3, 1).contiguous()
+        )
+
+        # Final mesh
+        m = mesh.Mesh(
+            n_vert,
+            load_mesh.t_pos_idx,
+            material={
+                'bsdf': cfg.bsdf,
+                'kd': kd_notex,
+                'ks': ready_specular,
+                'normal': ready_normal,
+            },
+            base=load_mesh # gets uvs etc from here
+        )
+
+        deformed_mesh_p3d = Meshes(verts=[m.v_pos], faces=[m.t_pos_idx])
+
+        render_mesh = create_scene([m.eval()], sz=512)
+        if it == 0:
+            base_mesh = render_mesh.clone()
+            base_mesh = mesh.auto_normals(base_mesh)
+            base_mesh = mesh.compute_tangents(base_mesh)
+        render_mesh = mesh.auto_normals(render_mesh)
+        render_mesh = mesh.compute_tangents(render_mesh)
+
+        if use_target_mesh:
+            # Target mesh
+            m_target = mesh.Mesh(
+                target_mesh.v_pos,
+                target_mesh.t_pos_idx,
+                material={
+                    'bsdf': cfg.bsdf,
+                    'kd': kd_notex,
+                    'ks': ready_specular,
+                    'normal': ready_normal,
+                },
+                base=target_mesh
+            )
+
+            render_target_mesh = create_scene([m_target.eval()], sz=512)
+            if it == 0:
+                base_target_mesh = render_target_mesh.clone()
+                base_target_mesh = mesh.auto_normals(base_target_mesh)
+                base_target_mesh = mesh.compute_tangents(base_target_mesh)
+            render_target_mesh = mesh.auto_normals(render_target_mesh)
+            render_target_mesh = mesh.compute_tangents(render_target_mesh)
+
+
+        # Logging mesh
+        if it % cfg.log_interval == 0:
+            with torch.no_grad():
+                params = get_camera_params(
+                    cfg.log_elev,
+                    rot_ang,
+                    cfg.log_dist,
+                    cfg.log_res,
+                    cfg.log_fov,
+                )
+                rot_ang += 5
+                log_mesh = mesh.unit_size(render_mesh.eval(params))
+                log_image = render.render_mesh(
+                    glctx,
+                    log_mesh,
+                    params['mvp'],
+                    params['campos'],
+                    params['lightpos'],
+                    cfg.log_light_power,
+                    cfg.log_res,
+                    1,
+                    background=torch.ones(1, cfg.log_res, cfg.log_res, 3).to(device)
+                )
+
+                log_image = video.ready_image(log_image)
+                logger.add_mesh('predicted_mesh', vertices=log_mesh.v_pos.unsqueeze(0), faces=log_mesh.t_pos_idx.unsqueeze(0), global_step=it)
+
+        if cfg.adapt_dist and it > 0:
+            with torch.no_grad():
+                v_pos = m.v_pos.clone()
+                vmin = v_pos.amin(dim=0)
+                vmax = v_pos.amax(dim=0)
+                v_pos -= (vmin + vmax) / 2
+                mult = torch.cat([v_pos.amin(dim=0), v_pos.amax(dim=0)]).abs().amax().cpu()
+                cams.dataset.dist_min = cfg.dist_min * mult
+                cams.dataset.dist_max = cfg.dist_max * mult
+
+        params_camera = next(iter(cams))
+        for key in params_camera:
+            params_camera[key] = params_camera[key].to(device)
+
+        final_mesh = render_mesh.eval(params_camera)
+        train_render = render.render_mesh(
+            glctx,
+            final_mesh,
+            params_camera['mvp'],
+            params_camera['campos'],
+            params_camera['lightpos'],
+            cfg.light_power,
+            cfg.train_res,
+            spp=1,
+            num_layers=1,
+            msaa=False,
+            background=params_camera['bkgs']
+        ).permute(0, 3, 1, 2)
+        train_render = resize(train_render, out_shape=(224, 224), interp_method=resize_method)
+
+        if use_target_mesh:
+            final_target_mesh = render_target_mesh.eval(params_camera)
+            train_target_render = render.render_mesh(
+                glctx,
+                final_target_mesh,
+                params_camera['mvp'],
+                params_camera['campos'],
+                params_camera['lightpos'],
+                cfg.light_power,
+                cfg.train_res,
+                spp=1,
+                num_layers=1,
+                msaa=False,
+                background=params_camera['bkgs']
+            ).permute(0, 3, 1, 2)
+            train_target_render = resize(train_target_render, out_shape=(224, 224), interp_method=resize_method)
+
+        train_rast_map = render.render_mesh(
+            glctx,
+            final_mesh,
+            params_camera['mvp'],
+            params_camera['campos'],
+            params_camera['lightpos'],
+            cfg.light_power,
+            cfg.train_res,
+            spp=1,
+            num_layers=1,
+            msaa=False,
+            background=params_camera['bkgs'],
+            return_rast_map=True
+        )
+
+        if it == 0:
+            params_camera = next(iter(cams))
+            for key in params_camera:
+                params_camera[key] = params_camera[key].to(device)
+        base_render = render.render_mesh(
+            glctx,
+            base_mesh.eval(params_camera),
+            params_camera['mvp'],
+            params_camera['campos'],
+            params_camera['lightpos'],
+            cfg.light_power,
+            cfg.train_res,
+            spp=1,
+            num_layers=1,
+            msaa=False,
+            background=params_camera['bkgs'],
+        ).permute(0, 3, 1, 2)
+        base_render = resize(base_render, out_shape=(224, 224), interp_method=resize_method)
+
+        if it % cfg.log_interval_im == 0:
+            log_idx = torch.randperm(cfg.batch_size)[:5]
+            s_log = train_render[log_idx, :, :, :]
+            s_log = torchvision.utils.make_grid(s_log)
+            ndarr = s_log.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
+            im = Image.fromarray(ndarr)
+            im.save(str(output_path / 'images' / f'epoch_{it}.png'))
+
+            if use_target_mesh:
+                s_log_target = train_target_render[log_idx, :, :, :]
+                s_log_target = torchvision.utils.make_grid(s_log_target)
+                ndarr = s_log_target.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
+                im = Image.fromarray(ndarr)
+                im.save(str(output_path / 'images' / f'epoch_{it}_target.png'))
+
+            obj.write_obj(
+                str(output_path / 'mesh_final'),
+                m.eval()
+            )
+
+        optimizer.zero_grad()
+
+
+        normalized_clip_render = (train_render - clip_mean[None, :, None, None]) / clip_std[None, :, None, None]
+
+        deformed_features = fclip.encode_image_tensors(train_render)
+        target_features = fclip.encode_image_tensors(train_target_render)
+        garment_loss = l1_avg(deformed_features, target_features)
+        l1_loss = l1_avg(train_render, train_target_render)
+
+        # We sample 10k points from the surface of each mesh
+        sample_src = sample_points_from_meshes(deformed_mesh_p3d, 10000)
+        sample_trg = sample_points_from_meshes(trg_mesh_p3d, 10000)
+
+        # We compare the two sets of pointclouds by computing (a) the chamfer loss
+        loss_chamfer, _ = chamfer_distance(sample_trg, sample_src)
+        loss_chamfer *= 25.
+        #
+        # and (b) the edge length of the predicted mesh
+        loss_edge = mesh_edge_loss(deformed_mesh_p3d)
+
+        # mesh normal consistency
+        loss_normal = mesh_normal_consistency(deformed_mesh_p3d)
+
+        # mesh laplacian smoothing
+        loss_laplacian = mesh_laplacian_smoothing(deformed_mesh_p3d, method="uniform")
+
+        loss_triangles = triangle_size_regularization(deformed_mesh_p3d.verts_list()[0])/100000.
+
+        logger.add_scalar('l1_loss', l1_loss, global_step=it)
+        logger.add_scalar('garment_loss', garment_loss, global_step=it)
+
+        # Jacobian regularization
+        r_loss = (((gt_jacobians) - torch.eye(3, 3, device=device)) ** 2).mean()
+        logger.add_scalar('jacobian_regularization', r_loss, global_step=it)
+
+        if cfg.consistency_loss_weight != 0:
+            consistency_loss = compute_mv_cl(final_mesh, fe, normalized_clip_render, params_camera, train_rast_map, cfg, device)
+        else:
+            consistency_loss = r_loss
+        logger.add_scalar('consistency_loss', consistency_loss, global_step=it)
+
+        logger.add_scalar('chamfer', loss_chamfer, global_step=it)
+        logger.add_scalar('edge', loss_edge, global_step=it)
+        logger.add_scalar('normal', loss_normal, global_step=it)
+        logger.add_scalar('laplacian', loss_laplacian, global_step=it)
+        logger.add_scalar('triangles', loss_triangles, global_step=it)
+
+
+        if it > 1000 and clip_flag:
+            cfg.clip_weight = 0
+            cfg.consistency_loss_weight = 0
+            cfg.regularize_jacobians_weight = 0.025
+            clip_flag = False
+        regularizers = loss_chamfer + loss_edge + loss_normal + loss_laplacian + loss_triangles
+        total_loss = (cfg.clip_weight * garment_loss + cfg.delta_clip_weight * l1_loss +
+                      cfg.regularize_jacobians_weight * r_loss + cfg.consistency_loss_weight * consistency_loss + regularizers)
+
+        logger.add_scalar('total_loss', total_loss, global_step=it)
+
+        total_loss.backward()
+        optimizer.step()
+        t_loop.set_description(
+                               f'L1 = {cfg.delta_clip_weight * l1_loss.item()}, '
+                               f'CLIP = {cfg.clip_weight * garment_loss.item()}, '
+                               f'Jacb = {cfg.regularize_jacobians_weight * r_loss.item()}, '
+                               f'MVC = {cfg.consistency_loss_weight * consistency_loss.item()}, '
+                               f'Chamf = {loss_chamfer.item()}, '
+                               f'Edge = {loss_edge.item()}, '
+                               f'Normal = {loss_normal.item()}, '
+                               f'Lapl = {loss_laplacian.item()}, '
+                               f'Triang = {loss_triangles.item()}, '
+                               f'Total = {total_loss.item()}')#_target
+
+    video.close()
+    obj.write_obj(
+        str(output_path / 'mesh_final'),
+        m.eval()
+    )
+    
+    return

main.py

@@ -0,0 +1,93 @@
+import os
+import yaml
+import torch
+import random
+import argparse
+import numpy as np
+
+from loop import loop
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', help='Path to config file', type=str, default='./example_config.yml')
+    parser.add_argument('--output_path', help='Output directory (will be created)', type=str, default=argparse.SUPPRESS)
+    parser.add_argument('--gpu', help='GPU index', type=int, default=argparse.SUPPRESS)
+    parser.add_argument('--seed', help='Random seed', type=int, default=argparse.SUPPRESS)
+
+    # CLIP-related
+    parser.add_argument('--text_prompt', help='Target text prompt', type=str, default=argparse.SUPPRESS)
+    parser.add_argument('--base_text_prompt', help='Base text prompt describing input mesh', type=str, default=argparse.SUPPRESS)
+    parser.add_argument('--clip_model', help='CLIP Model for text comparison', type=str, default=argparse.SUPPRESS)
+    parser.add_argument('--consistency_clip_model', help='CLIP Model for consistency', type=str, default=argparse.SUPPRESS)
+    parser.add_argument('--consistency_vit_stride', help='New stride for ViT patch interpolation', type=int, default=argparse.SUPPRESS)
+    parser.add_argument('--consistency_vit_layer', help='Which layer to take ViT patch features from (0-11)', type=int, default=argparse.SUPPRESS)
+
+    # Mesh
+    parser.add_argument('--mesh', help='Path to input mesh', type=str, default=argparse.SUPPRESS)
+    parser.add_argument('--retriangulate', help='Use isotropic remeshing', type=int, default=argparse.SUPPRESS, choices=[0, 1])
+
+    # Render settings
+    parser.add_argument('--bsdf', help='Render technique', type=str, default=argparse.SUPPRESS, choices=['diffuse', 'pbr'])
+
+    # Hyper-parameters
+    parser.add_argument('--lr', help='Learning rate', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--epochs', help='Number of optimization steps', type=int, default=argparse.SUPPRESS)
+    parser.add_argument('--clip_weight', help='Weight for CLIP loss', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--delta_clip_weight', help='Wight for delta-CLIP loss', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--regularize_jacobians_weight', help='Weight for jacobian regularization', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--consistency_loss_weight', help='Weight for viewpoint consistency penalty', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--consistency_elev_filter', help='Elev. angle threshold for filtering out pairs of viewpoints for consistency loss', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--consistency_azim_filter', help='Azim. angle threshold for filtering out pairs of viewpoints for consistency loss', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--batch_size', help='Number of images rendered at the same time', type=int, default=argparse.SUPPRESS)
+    parser.add_argument('--train_res', help='Resolution of render before downscaling to CLIP size', type=int, default=argparse.SUPPRESS)
+    parser.add_argument('--resize_method', help='Image downsampling/upsampling  method', type=str, default=argparse.SUPPRESS, choices=['cubic', 'linear', 'lanczos2', 'lanczos3'])
+    ## Camera Parameters ##
+    parser.add_argument('--fov_min', help='Minimum camera field of view angle during renders', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--fov_max', help='Maximum camera field of view angle during renders', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--dist_mi  n', help= 'Minimum distance of camera from mesh during renders', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--dist_max', help='Maximum distance of camera from mesh during renders', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--light_power', help='Light intensity', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--elev_alpha', help='Alpha parameter for Beta distribution for elevation sampling', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--elev_beta', help='Beta parameter for Beta distribution for elevation sampling', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--elev_max', help='Maximum elevation anglez in degree', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--azim_min', help='Minimum azimuth angle in degree',  type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--azim_max', help='Maximum azimuth angle in degree', type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--aug_loc', help='Offset mesh from center of image?', type=int, default=argparse.SUPPRESS, choices=[0, 1])
+    parser.add_argument('--aug_light', help='Augment the direction of light around the camera', type=int, default=argparse.SUPPRESS, choices=[0, 1])
+    parser.add_argument('--aug_bkg', help='Augment the background', type=int, default=argparse.SUPPRESS, choices=[0, 1])
+    parser.add_argument('--adapt_dist', help='Adjust camera distance to account for scale of shape', type=int, default=argparse.SUPPRESS, choices=[0, 1])
+
+    # Logging
+    parser.add_argument('--log_inter val', help='Interval for logging, every X epochs',  type=int, default=argparse.SUPPRESS)
+    parser.add_argument('--log_interval_im', help='Interval for logging renders image, every X epochs',  type=int, default=argparse.SUPPRESS)
+    parser.add_argument('--log_elev', help='Logging elevation angle',  type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--log_fov', help='Logging field of view',  type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--log_dist', help='Logging distance from object',  type=float, default=argparse.SUPPRESS)
+    parser.add_argument('--log_res', help='Logging render resolution',  type=int, default=argparse.SUPPRESS)
+    parser.add_argument('--log_light_power', help='Light intensity for logging',  type=float, default=argparse.SUPPRESS)
+
+    args = parser.parse_args()
+    if args.config is not None:
+        with open(args.config, 'r') as f:
+            try:
+                cfg = yaml.safe_load(f)
+            except yaml.YAMLError as e:
+                print(e)
+    
+    for key in vars(args):
+        cfg[key] = vars(args)[key]
+
+    print(yaml.dump(cfg, default_flow_style=False))
+    random.seed(cfg['seed'])
+    os.environ['PYTHONHASHSEED'] = str(cfg['seed'])
+    np.random.seed(cfg['seed'])
+    torch.manual_seed(cfg['seed'])
+    torch.cuda.manual_seed(cfg['seed'])
+    torch.backends.cudnn.deterministic = True
+
+    loop(cfg)
+    print('Done')
+
+if __name__ == '__main__':
+    main()
+

meshes/longsleeve.mtl

@@ -0,0 +1,13 @@
+# Blender MTL File: 'UV_warpings.blend'
+# Material Count: 1
+
+newmtl Material1770410.001
+Ns 250.000000
+Ka 1.000000 1.000000 1.000000
+Kd 0.800000 0.800000 0.800000
+Ks 0.500000 0.500000 0.500000
+Ke 0.000000 0.000000 0.000000
+Ni 1.450000
+d 1.000000
+illum 2
+map_Kd UV.png

meshes_target/jacket_sdf_new.obj

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b224a72b3385ab25aeb86fc357d091db111fb97caf470235b881b4a8116645a
+size 27668787

nvdiffmodeling/LICENSE.txt

@@ -0,0 +1,97 @@
+Copyright (c) 2021, NVIDIA Corporation. All rights reserved.
+
+
+Nvidia Source Code License (1-Way Commercial)
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+    License, each Licensor grants to you a perpetual, worldwide,
+    non-exclusive, royalty-free, copyright license to reproduce,
+    prepare derivative works of, publicly display, publicly perform,
+    sublicense and distribute its Work and any resulting derivative
+    works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+    if (a) you do so under this License, (b) you include a complete
+    copy of this License with your distribution, and (c) you retain
+    without modification any copyright, patent, trademark, or
+    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+    terms apply to the use, reproduction, and distribution of your
+    derivative works of the Work ("Your Terms") only if (a) Your Terms
+    provide that the use limitation in Section 3.3 applies to your
+    derivative works, and (b) you identify the specific derivative
+    works that are subject to Your Terms. Notwithstanding Your Terms,
+    this License (including the redistribution requirements in Section
+    3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+    may be used or intended for use non-commercially. The Work or
+    derivative works thereof may be used or intended for use by Nvidia
+    or its affiliates commercially or non-commercially. As used herein,
+    "non-commercially" means for research or evaluation purposes only
+    and not for any direct or indirect monetary gain.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+    against any Licensor (including any claim, cross-claim or
+    counterclaim in a lawsuit) to enforce any patents that you allege
+    are infringed by any Work, then your rights under this License from
+    such Licensor (including the grant in Section 2.1) will terminate
+    immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+    Licensor's or its affiliates' names, logos, or trademarks, except
+    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+    rights under this License (including the grant in Section 2.1) will
+    terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

nvdiffmodeling/src/material.py

@@ -0,0 +1,149 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import numpy as np
+import torch
+
+from . import util
+from . import texture
+from . import mesh
+
+######################################################################################
+# .mtl material format loading / storing
+######################################################################################
+
+def load_mtl(fn, clear_ks=True):
+    import re
+    mtl_path = os.path.dirname(fn)
+
+    # Read file
+    with open(fn) as f:
+        lines = f.readlines()
+
+    # Parse materials
+    materials = []
+    for line in lines:
+        split_line = re.split(' +|\t+|\n+', line.strip())
+        prefix = split_line[0].lower()
+        data = split_line[1:]
+        if 'newmtl' in prefix:
+            material = {'name' : data[0]}
+            materials += [material]
+        elif materials:
+            if 'bsdf' in prefix or 'map_kd' in prefix or 'map_ks' in prefix or 'bump' in prefix:
+                material[prefix] = data[0]
+            else:
+                material[prefix] = torch.tensor(tuple(float(d) for d in data), dtype=torch.float32, device='cuda')
+
+    # Convert everything to textures. Our code expects 'kd' and 'ks' to be texture maps. So replace constants with 1x1 maps
+    for mat in materials:
+        if not 'bsdf' in mat:
+            mat['bsdf'] = 'pbr'
+
+        if 'map_kd' in mat:
+            mat['kd'] = texture.load_texture2D(os.path.join(mtl_path, mat['map_kd']))
+        else:
+            mat['kd'] = texture.Texture2D(mat['kd'])
+        
+        if 'map_ks' in mat:
+            mat['ks'] = texture.load_texture2D(os.path.join(mtl_path, mat['map_ks']), channels=3)
+        else:
+            mat['ks'] = texture.Texture2D(mat['ks'])
+
+        if 'bump' in mat:
+            mat['normal'] = texture.load_texture2D(os.path.join(mtl_path, mat['bump']), lambda_fn=lambda x: x * 2 - 1, channels=3)
+
+        # Convert Kd from sRGB to linear RGB
+        mat['kd'] = texture.srgb_to_rgb(mat['kd'])
+
+        if clear_ks:
+            # Override ORM occlusion (red) channel by zeros. We hijack this channel
+            for mip in mat['ks'].getMips():
+                mip[..., 0] = 0.0 
+
+    return materials
+
+def save_mtl(fn, material):
+    folder = os.path.dirname(fn)
+    with open(fn, "w") as f:
+        f.write('newmtl defaultMat\n')
+        if material is not None:
+            f.write('bsdf   %s\n' % material['bsdf'])
+            f.write('map_kd texture_kd.png\n')
+            texture.save_texture2D(os.path.join(folder, 'texture_kd.png'), texture.rgb_to_srgb(material['kd']))
+            f.write('map_ks texture_ks.png\n')
+            texture.save_texture2D(os.path.join(folder, 'texture_ks.png'), material['ks'])
+            f.write('bump texture_n.png\n')
+            texture.save_texture2D(os.path.join(folder, 'texture_n.png'), material['normal'], lambda_fn=lambda x:(x+1)*0.5)
+        else:
+            f.write('Kd 1 1 1\n')
+            f.write('Ks 0 0 0\n')
+            f.write('Ka 0 0 0\n')
+            f.write('Tf 1 1 1\n')
+            f.write('Ni 1\n')
+            f.write('Ns 0\n')
+
+######################################################################################
+# Merge multiple materials into a single uber-material
+######################################################################################
+
+def _upscale_replicate(x, full_res):
+    x = x.permute(0, 3, 1, 2)
+    x = torch.nn.functional.pad(x, (0, full_res[1] - x.shape[3], 0, full_res[0] - x.shape[2]), 'replicate')
+    return x.permute(0, 2, 3, 1).contiguous()
+
+def merge_materials(materials, texcoords, tfaces, mfaces):
+    assert len(materials) > 0
+    for mat in materials:
+        assert mat['bsdf'] == materials[0]['bsdf'], "All materials must have the same BSDF (uber shader)"
+        assert ('normal' in mat) is ('normal' in materials[0]), "All materials must have either normal map enabled or disabled"
+
+    uber_material = {
+        'name' : 'uber_material',
+        'bsdf' : materials[0]['bsdf'],
+    }
+
+    textures = ['kd', 'ks', 'normal']
+
+    # Find maximum texture resolution across all materials and textures
+    max_res = None
+    for mat in materials:
+        for tex in textures:
+            tex_res = np.array(mat[tex].getRes()) if tex in mat else np.array([1, 1])
+            max_res = np.maximum(max_res, tex_res) if max_res is not None else tex_res
+    
+    # Compute size of compund texture and round up to nearest PoT
+    full_res = 2**np.ceil(np.log2(max_res * np.array([1, len(materials)]))).astype(np.int)
+
+    # Normalize texture resolution across all materials & combine into a single large texture
+    for tex in textures:
+        if tex in materials[0]:
+            tex_data = torch.cat(tuple(util.scale_img_nhwc(mat[tex].data, tuple(max_res)) for mat in materials), dim=2) # Lay out all textures horizontally, NHWC so dim2 is x
+            tex_data = _upscale_replicate(tex_data, full_res)
+            uber_material[tex] = texture.Texture2D(tex_data)
+
+    # Compute scaling values for used / unused texture area
+    s_coeff = [full_res[0] / max_res[0], full_res[1] / max_res[1]]
+
+    # Recompute texture coordinates to cooincide with new composite texture
+    new_tverts = {}
+    new_tverts_data = []
+    for fi in range(len(tfaces)):
+        matIdx = mfaces[fi]
+        for vi in range(3):
+            ti = tfaces[fi][vi]
+            if not (ti in new_tverts):
+                new_tverts[ti] = {}
+            if not (matIdx in new_tverts[ti]): # create new vertex
+                new_tverts_data.append([(matIdx + texcoords[ti][0]) / s_coeff[1], texcoords[ti][1] / s_coeff[0]]) # Offset texture coodrinate (x direction) by material id & scale to local space. Note, texcoords are (u,v) but texture is stored (w,h) so the indexes swap here
+                new_tverts[ti][matIdx] = len(new_tverts_data) - 1
+            tfaces[fi][vi] = new_tverts[ti][matIdx] # reindex vertex
+
+    return uber_material, new_tverts_data, tfaces
+

nvdiffmodeling/src/mesh.py

@@ -0,0 +1,510 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import numpy as np
+import torch
+
+from . import util
+from . import texture
+
+######################################################################################
+# Base mesh class
+######################################################################################
+class Mesh:
+    def __init__(self, v_pos=None, t_pos_idx=None, v_nrm=None, t_nrm_idx=None, v_tex=None, t_tex_idx=None, v_tng=None, t_tng_idx=None, 
+    v_weights=None, bone_mtx=None, material=None, base=None):
+        self.v_pos = v_pos
+        self.v_weights = v_weights
+        self.v_nrm = v_nrm
+        self.v_tex = v_tex
+        self.v_tng = v_tng
+        self.t_pos_idx = t_pos_idx
+        self.t_nrm_idx = t_nrm_idx
+        self.t_tex_idx = t_tex_idx
+        self.t_tng_idx = t_tng_idx
+        self.material = material
+        self.bone_mtx = bone_mtx
+
+        if base is not None:
+            self.copy_none(base)
+
+    def copy_none(self, other):
+        if self.v_pos is None:
+            self.v_pos = other.v_pos
+        if self.v_weights is None:
+            self.v_weights = other.v_weights
+        if self.t_pos_idx is None:
+            self.t_pos_idx = other.t_pos_idx
+        if self.v_nrm is None:
+            self.v_nrm = other.v_nrm
+        if self.t_nrm_idx is None:
+            self.t_nrm_idx = other.t_nrm_idx
+        if self.v_tex is None:
+            self.v_tex = other.v_tex
+        if self.t_tex_idx is None:
+            self.t_tex_idx = other.t_tex_idx
+        if self.v_tng is None:
+            self.v_tng = other.v_tng
+        if self.t_tng_idx is None:
+            self.t_tng_idx = other.t_tng_idx
+        if self.material is None:
+            self.material = other.material
+        if self.bone_mtx is None:
+            self.bone_mtx = other.bone_mtx
+
+    def get_frames(self):
+        return self.bone_mtx.shape[0] if self.bone_mtx is not None else 1
+
+    def clone(self):
+        out = Mesh(base=self)
+        if out.v_pos is not None:
+            out.v_pos = out.v_pos.clone()
+        if out.v_weights is not None:
+            out.v_weights = out.v_weights.clone()
+        if out.t_pos_idx is not None:
+            out.t_pos_idx = out.t_pos_idx.clone()
+        if out.v_nrm is not None:
+            out.v_nrm = out.v_nrm.clone()
+        if out.t_nrm_idx is not None:
+            out.t_nrm_idx = out.t_nrm_idx.clone()
+        if out.v_tex is not None:
+            out.v_tex = out.v_tex.clone()
+        if out.t_tex_idx is not None:
+            out.t_tex_idx = out.t_tex_idx.clone()
+        if out.v_tng is not None:
+            out.v_tng = out.v_tng.clone()
+        if out.t_tng_idx is not None:
+            out.t_tng_idx = out.t_tng_idx.clone()
+        if out.bone_mtx is not None:
+            out.bone_mtx = out.bone_mtx.clone()
+        return out
+
+    def eval(self, params={}):
+        return self
+
+######################################################################################
+# Compute AABB
+######################################################################################
+def aabb(mesh):
+    return torch.min(mesh.v_pos, dim=0).values, torch.max(mesh.v_pos, dim=0).values
+
+######################################################################################
+# Align base mesh to reference mesh:move & rescale to match bounding boxes.
+######################################################################################
+def unit_size(mesh):
+    with torch.no_grad():
+        vmin, vmax = aabb(mesh)
+        scale = 2 / torch.max(vmax - vmin).item()
+        v_pos = mesh.v_pos - (vmax + vmin) / 2 # Center mesh on origin
+        v_pos = v_pos * scale                  # Rescale to unit size
+
+        return Mesh(v_pos, base=mesh)
+
+def resize_mesh(mesh):
+    scale = 0.03234645293868976
+    vmax = torch.tensor([ 32.9707, 159.2754,  16.8091], device='cuda:0')
+    vmin = torch.tensor([-28.7435,  97.4448, -18.4702], device='cuda:0')
+    with torch.no_grad():
+        v_pos = (mesh.v_pos/scale) + (vmax + vmin) / 2
+    return Mesh(v_pos, base=mesh)
+
+
+######################################################################################
+# Center & scale mesh for rendering
+#
+# TODO: It should be better to compute camera position from animated reference mesh 
+# instead of centering and scaling all meshes
+######################################################################################
+def center_by_reference(base_mesh, ref_aabb, scale):
+    center = (ref_aabb[0] + ref_aabb[1]) * 0.5
+    scale = scale / torch.max(ref_aabb[1] - ref_aabb[0]).item()
+    v_pos = (base_mesh.v_pos - center[None, ...]) * scale
+    return Mesh(v_pos, base=base_mesh)
+
+######################################################################################
+# Rescale base-mesh from NDC [-1, 1] space to same dimensions as reference mesh
+######################################################################################
+def align_with_reference(base_mesh, ref_mesh): # TODO: Fix normals?
+    class mesh_op_align:
+        def __init__(self, base_mesh, ref_mesh):
+            self.base_mesh = base_mesh
+            with torch.no_grad():
+                b_vmin, b_vmax = aabb(base_mesh.eval())
+                r_vmin, r_vmax = aabb(ref_mesh.eval())
+                b_size = (b_vmax - b_vmin)
+                self.offset = (r_vmax + r_vmin) / 2
+                self.scale = (r_vmax - r_vmin) / torch.where(b_size > 1e-6, b_size, torch.ones_like(b_size))
+
+        def eval(self, params={}):
+            base_mesh = self.base_mesh.eval(params)
+            v_pos = base_mesh.v_pos * self.scale[None, ...] + self.offset[None, ...]
+            return Mesh(v_pos, base=base_mesh)
+
+    return mesh_op_align(base_mesh, ref_mesh)
+
+######################################################################################
+# Skinning
+######################################################################################
+
+# Helper function to skin homogeneous vectors
+def _skin_hvec(bone_mtx, weights, attr):
+    attr_out = torch.matmul(attr[None, ...], bone_mtx) * torch.transpose(weights, 0, 1)[..., None]
+    return attr_out.sum(dim=0)[:, :3]
+
+def skinning(mesh):
+    class mesh_op_skinning:
+        def __init__(self, input):
+            self.input = input
+            
+            mesh = self.input.eval()
+            t_pos_idx = mesh.t_pos_idx.detach().cpu().numpy() 
+            if mesh.t_nrm_idx is not None:
+                self.nrm_remap = self._compute_remap(t_pos_idx, mesh.v_nrm.shape[0], mesh.t_nrm_idx.detach().cpu().numpy())
+            if mesh.t_tng_idx is not None:
+                self.tng_remap = self._compute_remap(t_pos_idx, mesh.v_tng.shape[0], mesh.t_tng_idx.detach().cpu().numpy())
+
+        # Compute an index list with corresponding vertex index for each normal/tangent. Vertices may have multiple normals/tangents, but not the other way around
+        def _compute_remap(self, t_pos_idx, n_attrs, t_attr_idx):
+            assert len(t_pos_idx) == len(t_attr_idx)
+
+            attr_vtx_idx = [None] * n_attrs
+            for ti in range(0, len(t_pos_idx)):
+                for vi in range(0, 3):
+                    assert attr_vtx_idx[t_attr_idx[ti][vi]] is None or attr_vtx_idx[t_attr_idx[ti][vi]] == t_pos_idx[ti][vi], "Trying to skin a mesh with shared normals (normal with 2 sets of skinning weights)"
+                    attr_vtx_idx[t_attr_idx[ti][vi]] = t_pos_idx[ti][vi]
+
+            return torch.tensor(attr_vtx_idx, dtype=torch.int64, device='cuda')
+
+        def eval(self, params={}):
+            imesh = self.input.eval(params)
+
+            if imesh.v_weights is None or imesh.bone_mtx is None:
+                return imesh
+
+            # Compute frame (assume looping animation). Note, bone_mtx is stored [Frame, Bone, ...]
+            t_idx = params['time'] if 'time' in params else 0
+            t_idx = (t_idx % imesh.bone_mtx.shape[0]) # Loop animation
+            bone_mtx    = imesh.bone_mtx[t_idx, ...]
+            bone_mtx_it = torch.transpose(torch.inverse(bone_mtx), -2, -1)
+
+            weights = imesh.v_weights
+            assert weights.shape[1] == bone_mtx.shape[0]
+
+            # Normalize weights
+            weights = torch.abs(weights) # TODO: This stabilizes training, but I don't know why. All weights are already clamped to >0
+            weights = weights / torch.sum(weights, dim=1, keepdim=True)
+
+            # Skin position
+            v_pos_out = _skin_hvec(bone_mtx, weights, util.to_hvec(imesh.v_pos, 1))
+            
+            # Skin normal
+            v_nrm_out = None
+            if imesh.v_nrm is not None:
+                v_nrm_out = _skin_hvec(bone_mtx_it, weights[self.nrm_remap, ...], util.to_hvec(imesh.v_nrm, 0))
+                v_nrm_out = util.safe_normalize(v_nrm_out)
+            
+            # Skin tangent
+            v_tng_out = None
+            if imesh.v_tng is not None:
+                v_tng_out = _skin_hvec(bone_mtx, weights[self.tng_remap, ...], util.to_hvec(imesh.v_tng, 0))
+                v_tng_out = util.safe_normalize(v_tng_out)
+
+            if torch.is_anomaly_enabled():
+                assert torch.all(torch.isfinite(v_pos_out))
+                assert v_nrm_out is None or torch.all(torch.isfinite(v_nrm_out))
+                assert v_tng_out is None or torch.all(torch.isfinite(v_tng_out))
+
+            return Mesh(v_pos=v_pos_out[:, :3], v_nrm=v_nrm_out, v_tng=v_tng_out, base=imesh)
+
+    return mesh_op_skinning(mesh)
+
+# Skinning helper functions
+def guess_weights(base_mesh, ref_mesh, N=10):
+    base_v_pos = base_mesh.v_pos.detach().cpu().numpy()
+    ref_v_pos = ref_mesh.v_pos.detach().cpu().numpy()
+    ref_v_weights = ref_mesh.v_weights.detach().cpu().numpy()
+    base_v_weights = np.zeros((base_v_pos.shape[0], ref_v_weights.shape[1]), dtype=np.float32)
+    
+    for v_idx, vtx in enumerate(base_v_pos):
+        # Compute distance from current vertex to vertices in ref_mesh
+        diff = ref_v_pos - vtx[None, ...]
+        dist = np.sum(diff * diff, axis=-1)
+        idxs = np.argpartition(dist, N)        
+
+        # Get the N nearest vertices
+        sum_w = 0.0
+        sum_vtx_w = np.zeros_like(ref_v_weights[0,...])
+        for i in idxs[:N]:
+            sum_w += 1.0 / max(dist[i], 0.001)
+            sum_vtx_w += ref_v_weights[i, ...] / max(dist[i], 0.001)
+        base_v_weights[v_idx, ...] = sum_vtx_w / sum_w
+    
+    return base_v_weights
+
+def random_weights(base_mesh, ref_mesh):
+    init = np.random.uniform(size=(base_mesh.v_pos.shape[0], ref_mesh.v_weights.shape[1]), low=0.0, high=1.0)
+    return init / np.sum(init, axis=1, keepdims=True)
+
+
+######################################################################################
+# Simple smooth vertex normal computation
+######################################################################################
+def auto_normals(mesh):
+    class mesh_op_auto_normals:
+        def __init__(self, input):
+            self.input = input
+
+        def eval(self, params={}):
+            imesh = self.input.eval(params)
+
+            i0 = imesh.t_pos_idx[:, 0]
+            i1 = imesh.t_pos_idx[:, 1]
+            i2 = imesh.t_pos_idx[:, 2]
+
+            v0 = imesh.v_pos[i0, :]
+            v1 = imesh.v_pos[i1, :]
+            v2 = imesh.v_pos[i2, :]
+
+            face_normals = torch.cross(v1 - v0, v2 - v0)
+
+            # Splat face normals to vertices
+            v_nrm = torch.zeros_like(imesh.v_pos)
+            v_nrm.scatter_add_(0, i0[:, None].repeat(1,3), face_normals)
+            v_nrm.scatter_add_(0, i1[:, None].repeat(1,3), face_normals)
+            v_nrm.scatter_add_(0, i2[:, None].repeat(1,3), face_normals)
+
+            # Normalize, replace zero (degenerated) normals with some default value
+            v_nrm = torch.where(util.dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device='cuda'))
+
+            self.v_nrm = util.safe_normalize(v_nrm)
+
+            if torch.is_anomaly_enabled():
+                assert torch.all(torch.isfinite(self.v_nrm))
+
+            return Mesh(v_nrm = self.v_nrm, t_nrm_idx=imesh.t_pos_idx, base = imesh)
+
+    return mesh_op_auto_normals(mesh)
+
+######################################################################################
+# Compute tangent space from texture map coordinates
+# Follows http://www.mikktspace.com/ conventions
+######################################################################################
+def compute_tangents(mesh):
+    class mesh_op_compute_tangents:
+        def __init__(self, input):
+            self.input = input
+
+        def eval(self, params={}):
+            imesh = self.input.eval(params)
+
+            vn_idx = [None] * 3
+            pos = [None] * 3
+            tex = [None] * 3
+            for i in range(0,3):
+                pos[i] = imesh.v_pos[imesh.t_pos_idx[:, i]]
+                tex[i] = imesh.v_tex[imesh.t_tex_idx[:, i]]
+                vn_idx[i] = imesh.t_nrm_idx[:, i]
+
+            tangents = torch.zeros_like(imesh.v_nrm)
+            tansum   = torch.zeros_like(imesh.v_nrm)
+
+            # Compute tangent space for each triangle
+            uve1 = tex[1] - tex[0]
+            uve2 = tex[2] - tex[0]
+            pe1  = pos[1] - pos[0]
+            pe2  = pos[2] - pos[0]
+            
+            nom   = (pe1 * uve2[..., 1:2] - pe2 * uve1[..., 1:2])
+            denom = (uve1[..., 0:1] * uve2[..., 1:2] - uve1[..., 1:2] * uve2[..., 0:1])
+            
+            # Avoid division by zero for degenerated texture coordinates
+            tang = nom / torch.where(denom > 0.0, torch.clamp(denom, min=1e-6), torch.clamp(denom, max=-1e-6))
+
+            # Update all 3 vertices
+            for i in range(0,3):
+                idx = vn_idx[i][:, None].repeat(1,3)
+                tangents.scatter_add_(0, idx, tang)                # tangents[n_i] = tangents[n_i] + tang
+                tansum.scatter_add_(0, idx, torch.ones_like(tang)) # tansum[n_i] = tansum[n_i] + 1
+            tangents = tangents / tansum
+
+            # Normalize and make sure tangent is perpendicular to normal
+            tangents = util.safe_normalize(tangents)
+            tangents = util.safe_normalize(tangents - util.dot(tangents, imesh.v_nrm) * imesh.v_nrm)
+
+            self.v_tng = tangents
+
+            if torch.is_anomaly_enabled():
+                assert torch.all(torch.isfinite(tangents))
+
+            return Mesh(v_tng=self.v_tng, t_tng_idx=imesh.t_nrm_idx, base=imesh)
+
+    return mesh_op_compute_tangents(mesh)
+
+######################################################################################
+# Subdivide each triangle into 4 new ones. Edge midpoint subdivision
+######################################################################################
+def subdivide(mesh, steps=1):
+    class mesh_op_subdivide:
+        def __init__(self, input):
+            self.input = input
+            self.new_vtx_idx = [None] * 4
+            self.new_tri_idx = [None] * 4
+
+            imesh = self.input.eval()
+            v_attr = v_attr_orig = [imesh.v_pos, imesh.v_nrm, imesh.v_tex, imesh.v_tng]
+            v_idx = v_idx_orig = [imesh.t_pos_idx, imesh.t_nrm_idx, imesh.t_tex_idx, imesh.t_tng_idx]
+            
+            for i, attr in enumerate(v_attr):
+                if attr is not None:
+                    tri_idx = v_idx[i].cpu().numpy()
+
+                    # Find unique edges
+                    edge_fetch_a = []
+                    edge_fetch_b = []
+                    edge_verts = {}
+                    for tri in tri_idx:
+                        for e_idx in range(0, 3):
+                            v0 = tri[e_idx]
+                            v1 = tri[(e_idx + 1) % 3]
+                            if (v1, v0) not in edge_verts.keys():
+                                edge_verts[(v0, v1)] = [len(edge_fetch_a), v0, v1]
+                                edge_fetch_a += [v0]
+                                edge_fetch_b += [v1]
+
+                    # Create vertex fetch lists for computing midpoint vertices
+                    self.new_vtx_idx[i] = [torch.tensor(edge_fetch_a, dtype=torch.int64, device='cuda'), torch.tensor(edge_fetch_b, dtype=torch.int64, device='cuda')]
+
+                    # Create subdivided triangles
+                    new_tri_idx = []
+                    for tri in tri_idx:
+                        v0, v1, v2= tri
+                        h0 = (edge_verts[(v0, v1)][0] if (v0, v1) in edge_verts.keys() else edge_verts[(v1, v0)][0]) + attr.shape[0]
+                        h1 = (edge_verts[(v1, v2)][0] if (v1, v2) in edge_verts.keys() else edge_verts[(v2, v1)][0]) + attr.shape[0]
+                        h2 = (edge_verts[(v2, v0)][0] if (v2, v0) in edge_verts.keys() else edge_verts[(v0, v2)][0]) + attr.shape[0]
+                        new_tri_idx += [[v0, h0, h2], [h0, v1, h1], [h1, v2, h2], [h0, h1, h2]]
+                    self.new_tri_idx[i] = torch.tensor(new_tri_idx, dtype=torch.int64, device='cuda')
+
+        def eval(self, params={}):
+            imesh = self.input.eval(params)
+
+            v_attr = v_attr_orig = [imesh.v_pos, imesh.v_nrm, imesh.v_tex, imesh.v_tng]
+            v_idx = v_idx_orig = [imesh.t_pos_idx, imesh.t_nrm_idx, imesh.t_tex_idx, imesh.t_tng_idx]
+
+            for i, attr in enumerate(v_attr):
+                if attr is not None:
+                    # Create new edge midpoint attributes
+                    edge_attr = (attr[self.new_vtx_idx[i][0], :] + attr[self.new_vtx_idx[i][1], :]) * 0.5
+                    v_attr[i] = torch.cat([attr, edge_attr], dim=0)
+                    
+                    # Copy new triangle lists
+                    v_idx[i] = self.new_tri_idx[i]
+
+            return Mesh(v_attr[0], v_idx[0], v_attr[1], v_idx[1], v_attr[2], v_idx[2], v_attr[3], v_idx[3], base=imesh)
+
+    x = mesh
+    for i in range(steps):
+        x = mesh_op_subdivide(x)
+
+    bm = mesh.eval()
+    sm = x.eval()
+    v_attr_orig = [bm.v_pos, bm.v_nrm, bm.v_tex, bm.v_tng]
+    v_attr = [sm.v_pos, sm.v_nrm, sm.v_tex, sm.v_tng]
+    v_idx_orig = [bm.t_pos_idx, bm.t_nrm_idx, bm.t_tex_idx, bm.t_tng_idx]
+    v_idx = [sm.t_pos_idx, sm.t_nrm_idx, sm.t_tex_idx, sm.t_tng_idx]
+    print("Subdivided mesh:")
+    print("    Attrs:   [%6d, %6d, %6d, %6d] -> [%6d, %6d, %6d, %6d]" % tuple(list((a.shape[0] if a is not None else 0) for a in v_attr_orig) + list((a.shape[0] if a is not None else 0) for a in v_attr)))
+    print("    Indices: [%6d, %6d, %6d, %6d] -> [%6d, %6d, %6d, %6d]" % tuple(list((a.shape[0] if a is not None else 0) for a in v_idx_orig) + list((a.shape[0] if a is not None else 0) for a in v_idx)))
+
+    return x
+
+######################################################################################
+# Displacement mapping
+######################################################################################
+def displace(mesh, displacement_map, scale=1.0, keep_connectivity=True):
+    class mesh_op_displace:
+        def __init__(self, input, displacement_map, scale, keep_connectivity):
+            self.input = input
+            self.displacement_map = displacement_map
+            self.scale = scale
+            self.keep_connectivity = keep_connectivity
+        
+        def eval(self, params={}):
+            imesh = self.input.eval(params)
+
+            if self.keep_connectivity:
+                vd   = torch.zeros_like(imesh.v_pos)
+                vd_n = torch.zeros_like(imesh.v_pos)
+                for i in range(0, 3):
+                    v = imesh.v_pos[imesh.t_pos_idx[:, i], :]
+                    n = imesh.v_nrm[imesh.t_nrm_idx[:, i], :]
+                    t = imesh.v_tex[imesh.t_tex_idx[:, i], :]
+                    v_displ = v + n * self.scale * util.tex_2d(self.displacement_map, t)
+            
+                    splat_idx = imesh.t_pos_idx[:, i, None].repeat(1,3)
+                    vd.scatter_add_(0, splat_idx, v_displ)
+                    vd_n.scatter_add_(0, splat_idx, torch.ones_like(v_displ))
+
+                return Mesh(vd / vd_n, base=imesh)
+            else:
+                vd   = torch.zeros([imesh.v_tex.shape[0], 3], dtype=torch.float32, device='cuda')
+                vd_n = torch.zeros([imesh.v_tex.shape[0], 3], dtype=torch.float32, device='cuda')
+                for i in range(0, 3):
+                    v = imesh.v_pos[imesh.t_pos_idx[:, i], :]
+                    n = imesh.v_nrm[imesh.t_nrm_idx[:, i], :]
+                    t = imesh.v_tex[imesh.t_tex_idx[:, i], :]
+                    v_displ = v + n * self.scale * util.tex_2d(self.displacement_map, t)
+            
+                    splat_idx = imesh.t_tex_idx[:, i, None].repeat(1, 3)
+                    vd.scatter_add_(0, splat_idx, v_displ)
+                    vd_n.scatter_add_(0, splat_idx, torch.ones_like(v_displ))
+
+                return Mesh(vd / vd_n, mesh.t_tex_idx, base=imesh)
+
+    return mesh_op_displace(mesh, displacement_map, scale, keep_connectivity)
+
+
+######################################################################################
+# Utilities to merge meshes / materials. No mesh-ops or differentiable stuff here.
+######################################################################################
+
+def merge(mesh_a, mesh_b):
+    def _merge_attr_idx(a, b, a_idx, b_idx):
+        if a is None and b is None:
+            return None, None
+        elif a is not None and b is None:
+            return a, a_idx
+        elif a is None and b is not None:
+            return b, b_idx
+        else:
+            return torch.cat((a, b), dim=0), torch.cat((a_idx, b_idx + a.shape[0]), dim=0)
+
+    v_pos, t_pos_idx = _merge_attr_idx(mesh_a.v_pos, mesh_b.v_pos, mesh_a.t_pos_idx, mesh_b.t_pos_idx)
+    v_nrm, t_nrm_idx = _merge_attr_idx(mesh_a.v_nrm, mesh_b.v_nrm, mesh_a.t_nrm_idx, mesh_b.t_nrm_idx)
+    v_tng, t_tng_idx = _merge_attr_idx(mesh_a.v_tng, mesh_b.v_tng, mesh_a.t_tng_idx, mesh_b.t_tng_idx)
+    v_tex, t_tex_idx = _merge_attr_idx(mesh_a.v_tex, mesh_b.v_tex, mesh_a.t_tex_idx, mesh_b.t_tex_idx)
+
+    if mesh_a.v_weights is None and mesh_b.v_weights is None:
+        v_weights, bone_mtx = None, None
+    elif mesh_a.v_weights is not None and mesh_b.v_weights is None:
+        v_weights, bone_mtx = mesh_a.v_weights, mesh_a.bone_mtx
+    elif mesh_a.v_weights is None and mesh_b.v_weights is not None:
+        v_weights, bone_mtx = mesh_b.v_weights, mesh_b.bone_mtx
+    else:
+        if torch.all(mesh_a.bone_mtx == mesh_b.bone_mtx): # TODO: Wanted to test if same pointer
+            bone_mtx = mesh_a.bone_mtx
+            v_weights = torch.cat((mesh_a.v_weights, mesh_b.v_weights), dim=0)
+        else:
+            bone_mtx = torch.cat((mesh_a.bone_mtx, mesh_b.bone_mtx), dim=1) # Frame, Bone, ...
+        
+            # Weights need to be increased to account for all bones
+            v_wa = torch.nn.functional.pad(mesh_a.v_weights, [0, mesh_b.v_weights.shape[1]]) #Pad weights_a with shape of weights_b
+            v_wb = torch.nn.functional.pad(mesh_b.v_weights, [mesh_a.v_weights.shape[1], 0]) #Pad weights_b with shape of weights_a
+            v_weights = torch.cat((v_wa, v_wb), dim=0)
+
+    return Mesh(v_pos=v_pos, t_pos_idx=t_pos_idx, v_nrm=v_nrm, t_nrm_idx=t_nrm_idx, v_tng=v_tng, t_tng_idx=t_tng_idx, v_tex=v_tex, t_tex_idx=t_tex_idx, v_weights=v_weights, bone_mtx=bone_mtx, base=mesh_a)

nvdiffmodeling/src/obj.py

@@ -0,0 +1,215 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import numpy as np
+import torch
+
+from . import util
+from . import texture
+from . import mesh
+from . import material
+
+######################################################################################
+# Utility functions
+######################################################################################
+
+def _write_weights(folder, mesh):
+    if mesh.v_weights is not None:
+        file = os.path.join(folder, 'mesh.weights')
+        np.save(file, mesh.v_weights.detach().cpu().numpy())
+
+def _write_bones(folder, mesh):
+    if mesh.bone_mtx is not None:
+        file = os.path.join(folder, 'mesh.bones')
+        np.save(file, mesh.bone_mtx.detach().cpu().numpy())
+
+def _find_mat(materials, name):
+    for mat in materials:
+        if mat['name'] == name:
+            return mat
+    return materials[0] # Materials 0 is the default
+
+######################################################################################
+# Create mesh object from objfile
+######################################################################################
+
+def load_obj(filename, clear_ks=True, mtl_override=None):
+    obj_path = os.path.dirname(filename)
+
+    # Read entire file
+    with open(filename) as f:
+        lines = f.readlines()
+
+    # Load materials
+    all_materials = [
+        {
+            'name' : '_default_mat',
+            'bsdf' : 'falcor',
+            'kd'   : texture.Texture2D(torch.tensor([0.5, 0.5, 0.5], dtype=torch.float32, device='cuda')),
+            'ks'   : texture.Texture2D(torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32, device='cuda'))
+        }
+    ]
+    if mtl_override is None: 
+        for line in lines:
+            if len(line.split()) == 0:
+                continue
+            if line.split()[0] == 'mtllib':
+                all_materials += material.load_mtl(obj_path + os.path.join(line.split()[1]), clear_ks) # Read in entire material library #obj_path
+    else:
+        all_materials += material.load_mtl(mtl_override)
+
+    # load vertices
+    vertices, texcoords, normals  = [], [], []
+    for line in lines:
+        if len(line.split()) == 0:
+            continue
+        
+        prefix = line.split()[0].lower()
+        if prefix == 'v':
+            vertices.append([float(v) for v in line.split()[1:]][:3])
+        elif prefix == 'vt':
+            val = [float(v) for v in line.split()[1:]]
+            texcoords.append([val[0], 1.0 - val[1]])
+        elif prefix == 'vn':
+            normals.append([float(v) for v in line.split()[1:]])
+
+    # load faces
+    activeMatIdx = None
+    used_materials = []
+    faces, tfaces, nfaces, mfaces = [], [], [], []
+    for line in lines:
+        if len(line.split()) == 0:
+            continue
+
+        prefix = line.split()[0].lower()
+        if prefix == 'usemtl': # Track used materials
+            mat = _find_mat(all_materials, line.split()[1])
+            if not mat in used_materials:
+                used_materials.append(mat)
+            activeMatIdx = used_materials.index(mat)
+        elif prefix == 'f': # Parse face
+            vs = line.split()[1:]
+            nv = len(vs)
+            vv = vs[0].split('/')
+            v0 = int(vv[0]) - 1
+            if len(vv) > 1:
+                t0 = int(vv[1]) - 1 if vv[1] != "" else -1
+                n0 = int(vv[2]) - 1 if vv[2] != "" else -1
+            else:
+                t0 = -1
+                n0 = -1
+            for i in range(nv - 2): # Triangulate polygons
+                vv = vs[i + 1].split('/')
+                v1 = int(vv[0]) - 1
+                if len(vv) > 1:
+                    t1 = int(vv[1]) - 1 if vv[1] != "" else -1
+                    n1 = int(vv[2]) - 1 if vv[2] != "" else -1
+                else:
+                    t1 = -1
+                    n1 = -1
+                vv = vs[i + 2].split('/')
+                v2 = int(vv[0]) - 1
+                if len(vv) > 1:
+                    t2 = int(vv[1]) - 1 if vv[1] != "" else -1
+                    n2 = int(vv[2]) - 1 if vv[2] != "" else -1
+                else:
+                    t2 = -1
+                    n2 = -1
+                mfaces.append(activeMatIdx)
+                faces.append([v0, v1, v2])
+                tfaces.append([t0, t1, t2])
+                nfaces.append([n0, n1, n2])
+    assert len(tfaces) == len(faces) and len(nfaces) == len (faces)
+
+    # Create an "uber" material by combining all textures into a larger texture
+    if len(used_materials) > 1:
+        uber_material, texcoords, tfaces = material.merge_materials(used_materials, texcoords, tfaces, mfaces)
+    elif len(used_materials) == 1:
+        uber_material = used_materials[0]
+    else:
+        uber_material = None
+
+    vertices = torch.tensor(vertices, dtype=torch.float32, device='cuda')
+    texcoords = torch.tensor(texcoords, dtype=torch.float32, device='cuda') if len(texcoords) > 0 else None
+    normals = torch.tensor(normals, dtype=torch.float32, device='cuda') if len(normals) > 0 else None
+    
+    faces = torch.tensor(faces, dtype=torch.int64, device='cuda')
+    tfaces = torch.tensor(tfaces, dtype=torch.int64, device='cuda') if texcoords is not None else None
+    nfaces = torch.tensor(nfaces, dtype=torch.int64, device='cuda') if normals is not None else None
+
+    # Read weights and bones if available
+    try:
+        v_weights = torch.tensor(np.load(os.path.splitext(filename)[0] + ".weights.npy"), dtype=torch.float32, device='cuda')
+        bone_mtx = torch.tensor(np.load(os.path.splitext(filename)[0] + ".bones.npy"), dtype=torch.float32, device='cuda')
+    except:
+        v_weights, bone_mtx = None, None
+
+    return mesh.Mesh(vertices, faces, normals, nfaces, texcoords, tfaces, v_weights=v_weights, bone_mtx=bone_mtx, material=uber_material)
+
+######################################################################################
+# Save mesh object to objfile
+######################################################################################
+
+def write_obj(folder, mesh, verbose=True):
+    obj_file = os.path.join(folder, 'mesh.obj')
+    if verbose:
+        print("Writing mesh: ", obj_file)
+    with open(obj_file, "w") as f:
+        f.write("mtllib mesh.mtl\n")
+        f.write("g default\n")
+
+        v_pos = mesh.v_pos.detach().cpu().numpy() if mesh.v_pos is not None else None
+        v_nrm = mesh.v_nrm.detach().cpu().numpy() if mesh.v_nrm is not None else None
+        v_tex = mesh.v_tex.detach().cpu().numpy() if mesh.v_tex is not None else None
+
+        t_pos_idx = mesh.t_pos_idx.detach().cpu().numpy() if mesh.t_pos_idx is not None else None
+        t_nrm_idx = mesh.t_nrm_idx.detach().cpu().numpy() if mesh.t_nrm_idx is not None else None
+        t_tex_idx = mesh.t_tex_idx.detach().cpu().numpy() if mesh.t_tex_idx is not None else None
+        if verbose:
+            print("    writing %d vertices" % len(v_pos))
+        for v in v_pos:
+            f.write('v {} {} {} \n'.format(v[0], v[1], v[2]))
+       
+        if v_tex is not None:
+            if verbose:
+                print("    writing %d texcoords" % len(v_tex))
+            assert(len(t_pos_idx) == len(t_tex_idx))
+            for v in v_tex:
+                f.write('vt {} {} \n'.format(v[0], 1.0 - v[1]))
+
+        if v_nrm is not None:
+            if verbose:
+                print("    writing %d normals" % len(v_nrm))
+            assert(len(t_pos_idx) == len(t_nrm_idx))
+            for v in v_nrm:
+                f.write('vn {} {} {}\n'.format(v[0], v[1], v[2]))
+
+        # faces
+        f.write("s 1 \n")
+        f.write("g pMesh1\n")
+        f.write("usemtl defaultMat\n")
+
+        # Write faces
+        if verbose:
+            print("    writing %d faces" % len(t_pos_idx))
+        for i in range(len(t_pos_idx)):
+            f.write("f ")
+            for j in range(3):
+                f.write(' %s/%s/%s' % (str(t_pos_idx[i][j]+1), '' if v_tex is None else str(t_tex_idx[i][j]+1), '' if v_nrm is None else str(t_nrm_idx[i][j]+1)))
+            f.write("\n")
+
+    mtl_file = os.path.join(folder, 'mesh.mtl')
+    if verbose:
+        print("Writing material: ", mtl_file)
+    material.save_mtl(mtl_file, mesh.material)
+
+    _write_weights(folder, mesh)
+    _write_bones(folder, mesh)
+    if verbose:
+        print("Done exporting mesh")

nvdiffmodeling/src/regularizer.py

@@ -0,0 +1,197 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import numpy as np
+import torch
+
+from . import util
+from . import texture
+
+######################################################################################
+# Computes the avergage edge length of a mesh. 
+# Rough estimate of the tessellation of a mesh. Can be used e.g. to clamp gradients
+######################################################################################
+def avg_edge_length(opt_mesh):
+    with torch.no_grad():
+        opt_mesh = opt_mesh.eval()
+        nVerts = opt_mesh.v_pos.shape[0]
+        t_pos_idx = opt_mesh.t_pos_idx.detach().cpu().numpy() 
+
+        # Find unique edges
+        ix_i = []
+        ix_j = []
+        edge_verts = {}
+        for tri in t_pos_idx:
+            for (i0, i1) in [(tri[0], tri[1]), (tri[1], tri[2]), (tri[2], tri[0])]:
+                if (i1, i0) not in edge_verts.keys():
+                    edge_verts[(i0, i1)] = True
+                    ix_i += [i0]
+                    ix_j += [i1]
+
+        # Setup torch tensors
+        ix_i = torch.tensor(ix_i, dtype=torch.int64, device='cuda')
+        ix_j = torch.tensor(ix_j, dtype=torch.int64, device='cuda')
+
+        # Gather edge vertex pairs
+        x_i = opt_mesh.v_pos[ix_i, :]
+        x_j = opt_mesh.v_pos[ix_j, :]
+
+        # Compute edge length
+        term = torch.sqrt((x_j - x_i)**2)
+
+        # Compute avg edge length
+        return (torch.sum(term) / len(x_i)).item()
+
+######################################################################################
+# Edge length regularizer 
+######################################################################################
+def edge_length_regularizer(mesh):
+    class mesh_op_edge_length_regularizer:
+        def __init__(self, mesh):
+            self.mesh = mesh
+            
+            mesh = mesh.eval()
+            nVerts = mesh.v_pos.shape[0]
+            t_pos_idx = mesh.t_pos_idx.detach().cpu().numpy() 
+
+            # Find unique edges
+            ix_i = []
+            ix_j = []
+            edge_verts = {}
+            for tri in t_pos_idx:
+                for (i0, i1) in [(tri[0], tri[1]), (tri[1], tri[2]), (tri[2], tri[0])]:
+                    if (i1, i0) not in edge_verts.keys():
+                        edge_verts[(i0, i1)] = True
+                        ix_i += [i0]
+                        ix_j += [i1]
+
+            # Setup torch tensors
+            self.ix_i = torch.tensor(ix_i, dtype=torch.int64, device='cuda')
+            self.ix_j = torch.tensor(ix_j, dtype=torch.int64, device='cuda')
+
+        def eval(self, params={}):
+            mesh = self.mesh.eval(params)
+
+            # Gather edge vertex pairs
+            x_i = mesh.v_pos[self.ix_i, :]
+            x_j = mesh.v_pos[self.ix_j, :]
+
+            # Compute edge length
+            term = torch.sqrt((x_j - x_i)**2 + 1e-20)
+
+            # Compute avg edge length
+            return torch.var(term)
+
+    return mesh_op_edge_length_regularizer(mesh)
+
+######################################################################################
+# Laplacian regularization using umbrella operator (Fujiwara / Desbrun).
+# https://mgarland.org/class/geom04/material/smoothing.pdf
+######################################################################################
+def laplace_regularizer_const(opt_mesh, base_mesh=None):
+    class mesh_op_laplace_regularizer_const:
+        def __init__(self, opt_mesh, base_mesh):
+            self.inputs = [opt_mesh, base_mesh]
+
+            opt_mesh = opt_mesh.eval()
+            self.nVerts = opt_mesh.v_pos.shape[0]
+            t_pos_idx = opt_mesh.t_pos_idx.detach().cpu().numpy() 
+
+            # Build vertex neighbor rings
+            vtx_n = [[] for _ in range(self.nVerts)]
+            for tri in t_pos_idx:
+                for (i0, i1) in [(tri[0], tri[1]), (tri[1], tri[2]), (tri[2], tri[0])]:
+                    vtx_n[i0].append(i1)
+
+            # Collect index/weight pairs to compute each Laplacian vector for each vertex.
+            # Similar notation to https://mgarland.org/class/geom04/material/smoothing.pdf
+            ix_j, ix_i, w_ij = [], [], []
+            for i in range(self.nVerts):
+                m = len(vtx_n[i])
+                ix_i += [i] * m
+                ix_j += vtx_n[i]
+                w_ij += [1.0 / m] * m
+
+            # Setup torch tensors
+            self.ix_i = torch.tensor(ix_i, dtype=torch.int64, device='cuda')
+            self.ix_j = torch.tensor(ix_j, dtype=torch.int64, device='cuda')
+            self.w_ij = torch.tensor(w_ij, dtype=torch.float32, device='cuda')[:, None]
+
+        def eval(self, params={}):
+            opt_mesh = self.inputs[0].eval(params)
+            base_mesh = self.inputs[1].eval(params) if self.inputs[1] is not None else None
+
+            # differences or absolute version (see paper)
+            if base_mesh is not None:
+                v_pos = opt_mesh.v_pos - base_mesh.v_pos
+            else:
+                v_pos = opt_mesh.v_pos
+
+            # Gather edge vertex pairs
+            x_i = v_pos[self.ix_i, :]
+            x_j = v_pos[self.ix_j, :]
+
+            # Compute Laplacian differences: (x_j - x_i) * w_ij
+            term = (x_j - x_i) * self.w_ij
+
+            # Sum everyhing
+            term = util.segment_sum(term, self.ix_i)
+            
+            return torch.mean(term**2)
+    
+    return mesh_op_laplace_regularizer_const(opt_mesh, base_mesh)
+
+######################################################################################
+# Curvature based regularizer
+######################################################################################
+def face_normal_regularizer(opt_mesh):
+    class mesh_op_face_normal_regularizer:
+        def __init__(self, opt_mesh):
+            self.input = opt_mesh
+
+            imesh = opt_mesh.eval()
+            self.nVerts = imesh.v_pos.shape[0]
+            t_pos_idx = imesh.t_pos_idx.detach().cpu().numpy() 
+
+            # Generate edge lists
+            edge_tris = {}
+            for tri_idx, tri in enumerate(t_pos_idx):
+                for (i0, i1) in [(tri[0], tri[1]), (tri[1], tri[2]), (tri[2], tri[0])]:
+                    if (i1, i0) in edge_tris.keys():
+                        edge_tris[(i1, i0)] += [tri_idx]
+                    else:
+                        edge_tris[(i0, i1)] = [tri_idx]
+
+            # Get all good edges with 2 incident triangles
+            shared_edge_idx = []
+            for edge in edge_tris.values():
+                if len(edge) == 2:
+                    shared_edge_idx += [edge]
+            self.edge_tri_idx = torch.tensor(shared_edge_idx, dtype=torch.int64, device='cuda')
+
+        def eval(self, params={}):
+            imesh = self.input.eval(params)
+
+            # Compute face normals
+            v0 = imesh.v_pos[imesh.t_pos_idx[:, 0], :]
+            v1 = imesh.v_pos[imesh.t_pos_idx[:, 1], :]
+            v2 = imesh.v_pos[imesh.t_pos_idx[:, 2], :]
+            face_normals = util.safe_normalize(torch.cross(v1 - v0, v2 - v0))
+
+            # Fetch normals for both faces sharind an edge
+            n0 = face_normals[self.edge_tri_idx[:, 0], :]
+            n1 = face_normals[self.edge_tri_idx[:, 1], :]
+
+            # Compute error metric based on normal difference
+            term = torch.clamp(util.dot(n0, n1), min=-1.0, max=1.0)
+            term = (1.0 - term) * 0.5
+
+            return torch.mean(torch.abs(term))
+    
+    return mesh_op_face_normal_regularizer(opt_mesh)

nvdiffmodeling/src/render.py

@@ -0,0 +1,223 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import nvdiffrast.torch as dr
+
+from . import util
+from . import mesh
+from . import renderutils as ru
+
+# ==============================================================================================
+#  Helper functions
+# ==============================================================================================
+def interpolate(attr, rast, attr_idx, rast_db=None):
+    return dr.interpolate(attr.contiguous(), rast, attr_idx, rast_db=rast_db, diff_attrs=None if rast_db is None else 'all')
+
+# ==============================================================================================
+#  pixel shader
+# ==============================================================================================
+def shade(
+        gb_pos,
+        gb_geometric_normal,
+        gb_normal,
+        gb_tangent,
+        gb_texc,
+        gb_texc_deriv,
+        view_pos,
+        light_pos,
+        light_power,
+        material,
+        min_roughness
+    ):
+
+    ################################################################################
+    # Texture lookups
+    ################################################################################
+
+    kd = material['kd'].sample(gb_texc, gb_texc_deriv)
+    ks = material['ks'].sample(gb_texc, gb_texc_deriv)[..., 0:3] # skip alpha
+    perturbed_nrm = None
+    if 'normal' in material:
+        perturbed_nrm = material['normal'].sample(gb_texc, gb_texc_deriv)
+
+    gb_normal = ru.prepare_shading_normal(gb_pos, view_pos, perturbed_nrm, gb_normal, gb_tangent, gb_geometric_normal, two_sided_shading=True, opengl=True)
+
+    # Separate kd into alpha and color, default alpha = 1
+    alpha = kd[..., 3:4] if kd.shape[-1] == 4 else torch.ones_like(kd[..., 0:1]) 
+    kd = kd[..., 0:3]
+
+    ################################################################################
+    # Evaluate BSDF
+    ################################################################################
+
+    assert 'bsdf' in material, "Material must specify a BSDF type"
+    if material['bsdf'] == 'pbr':
+        shaded_col = ru.pbr_bsdf(kd, ks, gb_pos, gb_normal, view_pos, light_pos, min_roughness) * light_power
+    elif material['bsdf'] == 'diffuse':
+        shaded_col = kd * ru.lambert(gb_normal, util.safe_normalize(light_pos - gb_pos)) * light_power
+    elif material['bsdf'] == 'normal':
+        shaded_col = (gb_normal + 1.0)*0.5
+    elif material['bsdf'] == 'tangent':
+        shaded_col = (gb_tangent + 1.0)*0.5
+    else:
+        assert False, "Invalid BSDF '%s'" % material['bsdf']
+
+    out = torch.cat((shaded_col, alpha), dim=-1)
+
+    return out
+
+# ==============================================================================================
+#  Render a depth slice of the mesh (scene), some limitations:
+#  - Single mesh
+#  - Single light
+#  - Single material
+# ==============================================================================================
+def render_layer(
+        rast,
+        rast_deriv,
+        mesh,
+        view_pos,
+        light_pos,
+        light_power,
+        resolution,
+        min_roughness,
+        spp,
+        msaa
+    ):
+
+    full_res = resolution*spp
+
+    ################################################################################
+    # Rasterize
+    ################################################################################
+
+    # Scale down to shading resolution when MSAA is enabled, otherwise shade at full resolution
+    if spp > 1 and msaa:
+        rast_out_s = util.scale_img_nhwc(rast, [resolution, resolution], mag='nearest', min='nearest')
+        rast_out_deriv_s = util.scale_img_nhwc(rast_deriv, [resolution, resolution], mag='nearest', min='nearest') * spp
+    else:
+        rast_out_s = rast
+        rast_out_deriv_s = rast_deriv
+
+    ################################################################################
+    # Interpolate attributes
+    ################################################################################
+
+    # Interpolate world space position
+    gb_pos, _ = interpolate(mesh.v_pos[None, ...], rast_out_s, mesh.t_pos_idx.int())
+
+    # Compute geometric normals. We need those because of bent normals trick (for bump mapping)
+    v0 = mesh.v_pos[mesh.t_pos_idx[:, 0], :]
+    v1 = mesh.v_pos[mesh.t_pos_idx[:, 1], :]
+    v2 = mesh.v_pos[mesh.t_pos_idx[:, 2], :]
+    face_normals = util.safe_normalize(torch.cross(v1 - v0, v2 - v0))
+    face_normal_indices = (torch.arange(0, face_normals.shape[0], dtype=torch.int64, device='cuda')[:, None]).repeat(1, 3)
+    gb_geometric_normal, _ = interpolate(face_normals[None, ...], rast_out_s, face_normal_indices.int())
+
+    # Compute tangent space
+    assert mesh.v_nrm is not None and mesh.v_tng is not None
+    gb_normal, _ = interpolate(mesh.v_nrm[None, ...], rast_out_s, mesh.t_nrm_idx.int())
+    gb_tangent, _ = interpolate(mesh.v_tng[None, ...], rast_out_s, mesh.t_tng_idx.int()) # Interpolate tangents
+
+    # Texure coordinate
+    assert mesh.v_tex is not None
+    gb_texc, gb_texc_deriv = interpolate(mesh.v_tex[None, ...], rast_out_s, mesh.t_tex_idx.int(), rast_db=rast_out_deriv_s)
+
+    ################################################################################
+    # Shade
+    ################################################################################
+
+    color = shade(gb_pos, gb_geometric_normal, gb_normal, gb_tangent, gb_texc, gb_texc_deriv, 
+        view_pos, light_pos, light_power, mesh.material, min_roughness)
+
+    ################################################################################
+    # Prepare output
+    ################################################################################
+
+    # Scale back up to visibility resolution if using MSAA
+    if spp > 1 and msaa:
+        color = util.scale_img_nhwc(color, [full_res, full_res], mag='nearest', min='nearest')
+
+    # Return color & raster output for peeling
+    return color
+
+
+# ==============================================================================================
+#  Render a depth peeled mesh (scene), some limitations:
+#  - Single mesh
+#  - Single light
+#  - Single material
+# ==============================================================================================
+def render_mesh(
+        ctx,
+        mesh,
+        mtx_in,
+        view_pos,
+        light_pos,
+        light_power,
+        resolution,
+        spp                       = 1,
+        num_layers                = 1,
+        msaa                      = False,
+        background                = None,
+        antialias                 = True,
+        min_roughness             = 0.08,
+        return_rast_map           = False,
+    ):
+    assert not (return_rast_map and num_layers > 1)
+
+    def prepare_input_vector(x):
+        x = torch.tensor(x, dtype=torch.float32, device='cuda') if not torch.is_tensor(x) else x
+        return x[:, None, None, :] if len(x.shape) == 2 else x
+
+    full_res = resolution*spp
+
+    # Convert numpy arrays to torch tensors
+    mtx_in      = torch.tensor(mtx_in, dtype=torch.float32, device='cuda') if not torch.is_tensor(mtx_in) else mtx_in
+    light_pos   = prepare_input_vector(light_pos)
+    light_power = prepare_input_vector(light_power)
+    view_pos    = prepare_input_vector(view_pos)
+
+    # clip space transform
+    v_pos_clip = ru.xfm_points(mesh.v_pos[None, ...], mtx_in)
+
+    # Render all layers front-to-back
+    layers = []
+    with dr.DepthPeeler(ctx, v_pos_clip, mesh.t_pos_idx.int(), [resolution*spp, resolution*spp]) as peeler:
+        for _ in range(num_layers):
+            rast, db = peeler.rasterize_next_layer()
+            layers += [(render_layer(rast, db, mesh, view_pos, light_pos, light_power, resolution, min_roughness, spp, msaa), rast)]
+
+    if return_rast_map:
+        return rast.detach()
+
+    # Clear to background layer
+    if background is not None:
+        assert background.shape[1] == resolution and background.shape[2] == resolution
+        if spp > 1:
+            background = util.scale_img_nhwc(background, [full_res, full_res], mag='nearest', min='nearest')
+        accum_col = background
+    else:
+        accum_col = torch.zeros(size=(1, full_res, full_res, 3), dtype=torch.float32, device='cuda')
+
+    # Composite BACK-TO-FRONT
+    for color, rast in reversed(layers):
+        alpha     = (rast[..., -1:] > 0) * color[..., 3:4]
+        accum_col = torch.lerp(accum_col, color[..., 0:3], alpha)
+        if antialias:
+            accum_col = dr.antialias(accum_col.contiguous(), rast, v_pos_clip, mesh.t_pos_idx.int()) # TODO: need to support bfloat16
+
+    # Downscale to framebuffer resolution. Use avg pooling 
+    out = util.avg_pool_nhwc(accum_col, spp) if spp > 1 else accum_col
+
+    return out
+
+

nvdiffmodeling/src/renderutils/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from .ops import xfm_points, xfm_vectors, image_loss, prepare_shading_normal, lambert, pbr_specular, pbr_bsdf, _fresnel_shlick, _ndf_ggx, _lambda_ggx, _masking_smith
+__all__ = ["xfm_vectors", "xfm_points", "image_loss", "prepare_shading_normal", "lambert", "pbr_specular", "pbr_bsdf", "_fresnel_shlick", "_ndf_ggx", "_lambda_ggx", "_masking_smith", ]

nvdiffmodeling/src/renderutils/bsdf.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import math
+import torch
+
+NORMAL_THRESHOLD = 0.1
+
+################################################################################
+# Vector utility functions
+################################################################################
+
+def _dot(x, y):
+    return torch.sum(x*y, -1, keepdim=True)
+
+def _reflect(x, n):
+    return 2*_dot(x, n)*n - x
+
+def _safe_normalize(x):
+    return torch.nn.functional.normalize(x, dim = -1)
+
+def _bend_normal(view_vec, smooth_nrm, geom_nrm, two_sided_shading):
+    # Swap normal direction for backfacing surfaces
+    if two_sided_shading:
+        smooth_nrm = torch.where(_dot(geom_nrm, view_vec) > 0, smooth_nrm, -smooth_nrm)
+        geom_nrm   = torch.where(_dot(geom_nrm, view_vec) > 0, geom_nrm, -geom_nrm)
+
+    t = torch.clamp(_dot(view_vec, smooth_nrm) / NORMAL_THRESHOLD, min=0, max=1)
+    return torch.lerp(geom_nrm, smooth_nrm, t)
+
+
+def _perturb_normal(perturbed_nrm, smooth_nrm, smooth_tng, opengl):
+    smooth_bitang = _safe_normalize(torch.cross(smooth_tng, smooth_nrm))
+    if opengl:
+        shading_nrm = smooth_tng * perturbed_nrm[..., 0:1] - smooth_bitang * perturbed_nrm[..., 1:2] + smooth_nrm * torch.clamp(perturbed_nrm[..., 2:3], min=0.0)
+    else:
+        shading_nrm = smooth_tng * perturbed_nrm[..., 0:1] + smooth_bitang * perturbed_nrm[..., 1:2] + smooth_nrm * torch.clamp(perturbed_nrm[..., 2:3], min=0.0)
+    return _safe_normalize(shading_nrm)
+
+def bsdf_prepare_shading_normal(pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm, two_sided_shading, opengl):
+    smooth_nrm = _safe_normalize(smooth_nrm)
+    smooth_tng = _safe_normalize(smooth_tng)
+    view_vec   = _safe_normalize(view_pos - pos)
+    shading_nrm = _perturb_normal(perturbed_nrm, smooth_nrm, smooth_tng, opengl)
+    return _bend_normal(view_vec, shading_nrm, geom_nrm, two_sided_shading)
+
+################################################################################
+# Simple lambertian diffuse BSDF
+################################################################################
+
+def bsdf_lambert(nrm, wi):
+    return torch.clamp(_dot(nrm, wi), min=0.0) / math.pi
+
+################################################################################
+# Phong specular, loosely based on mitsuba implementation
+################################################################################
+
+def bsdf_phong(nrm, wo, wi, N):
+    dp_r = torch.clamp(_dot(_reflect(wo, nrm), wi), min=0.0, max=1.0)
+    dp_l = torch.clamp(_dot(nrm, wi), min=0.0, max=1.0)
+    return (dp_r ** N) * dp_l * (N + 2) / (2 * math.pi)
+
+################################################################################
+# PBR's implementation of GGX specular
+################################################################################
+
+specular_epsilon = 1e-4
+
+def bsdf_fresnel_shlick(f0, f90, cosTheta):
+    _cosTheta = torch.clamp(cosTheta, min=specular_epsilon, max=1.0 - specular_epsilon)
+    return f0 + (f90 - f0) * (1.0 - _cosTheta) ** 5.0
+
+def bsdf_ndf_ggx(alphaSqr, cosTheta):
+    _cosTheta = torch.clamp(cosTheta, min=specular_epsilon, max=1.0 - specular_epsilon)
+    d = (_cosTheta * alphaSqr - _cosTheta) * _cosTheta + 1
+    return alphaSqr / (d * d * math.pi)
+
+def bsdf_lambda_ggx(alphaSqr, cosTheta):
+    _cosTheta = torch.clamp(cosTheta, min=specular_epsilon, max=1.0 - specular_epsilon)
+    cosThetaSqr = _cosTheta * _cosTheta
+    tanThetaSqr = (1.0 - cosThetaSqr) / cosThetaSqr
+    res = 0.5 * (torch.sqrt(1 + alphaSqr * tanThetaSqr) - 1.0)
+    return res
+
+def bsdf_masking_smith_ggx_correlated(alphaSqr, cosThetaI, cosThetaO):
+    lambdaI = bsdf_lambda_ggx(alphaSqr, cosThetaI)
+    lambdaO = bsdf_lambda_ggx(alphaSqr, cosThetaO)
+    return 1 / (1 + lambdaI + lambdaO)
+
+def bsdf_pbr_specular(col, nrm, wo, wi, alpha, min_roughness=0.08):
+    _alpha = torch.clamp(alpha, min=min_roughness*min_roughness, max=1.0)
+    alphaSqr = _alpha * _alpha
+
+    h = _safe_normalize(wo + wi)
+    woDotN = _dot(wo, nrm)
+    wiDotN = _dot(wi, nrm)
+    woDotH = _dot(wo, h)
+    nDotH  = _dot(nrm, h)
+
+    D = bsdf_ndf_ggx(alphaSqr, nDotH)
+    G = bsdf_masking_smith_ggx_correlated(alphaSqr, woDotN, wiDotN)
+    F = bsdf_fresnel_shlick(col, 1, woDotH)
+
+    w = F * D * G * 0.25 / torch.clamp(woDotN, min=specular_epsilon)
+
+    frontfacing = (woDotN > specular_epsilon) & (wiDotN > specular_epsilon)
+    return torch.where(frontfacing, w, torch.zeros_like(w))
+
+def bsdf_pbr(kd, arm, pos, nrm, view_pos, light_pos, min_roughness=0.08):
+    wo = _safe_normalize(view_pos - pos)
+    wi = _safe_normalize(light_pos - pos)
+
+    spec_str  = arm[..., 0:1] # x component
+    roughness = arm[..., 1:2] # y component
+    metallic  = arm[..., 2:3] # z component
+    ks = (0.04 * (1.0 - metallic) + kd * metallic) * (1 - spec_str)
+    kd = kd * (1.0 - metallic)
+
+    diffuse = kd * bsdf_lambert(nrm, wi)
+    specular = bsdf_pbr_specular(ks, nrm, wo, wi, roughness*roughness, min_roughness=min_roughness)
+    return diffuse + specular

nvdiffmodeling/src/renderutils/c_src/bsdf.cu

@@ -0,0 +1,551 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include "common.h"
+#include "bsdf.h"
+
+#define SPECULAR_EPSILON 1e-4f
+
+//------------------------------------------------------------------------
+// Lambert functions
+
+__device__ float fwdLambert(const vec3f nrm, const vec3f wi)
+{
+    return max(dot(nrm, wi) / M_PI, 0.0f);
+}
+
+__device__ void bwdLambert(const vec3f nrm, const vec3f wi, vec3f& d_nrm, vec3f& d_wi, const float d_out)
+{
+    if (dot(nrm, wi) > 0.0f)
+        bwdDot(nrm, wi, d_nrm, d_wi, d_out / M_PI);
+}
+
+//------------------------------------------------------------------------
+// Fresnel Schlick 
+
+__device__ vec3f fwdFresnelSchlick(const vec3f f0, const vec3f f90, const float cosTheta)
+{
+    float _cosTheta = clamp(cosTheta, SPECULAR_EPSILON, 1.0f - SPECULAR_EPSILON);
+    float scale = powf(1.0f - _cosTheta, 5.0f);
+    return f0 * (1.0f - scale) + f90 * scale;
+}
+
+__device__ void bwdFresnelSchlick(const vec3f f0, const vec3f f90, const float cosTheta, vec3f& d_f0, vec3f& d_f90, float& d_cosTheta, const vec3f d_out)
+{
+    float _cosTheta = clamp(cosTheta, SPECULAR_EPSILON, 1.0f - SPECULAR_EPSILON);
+    float scale = pow(max(1.0f - _cosTheta, 0.0f), 5.0f);
+    d_f0 += d_out * (1.0 - scale);
+    d_f90 += d_out * scale;
+    if (cosTheta >= SPECULAR_EPSILON && cosTheta < 1.0f - SPECULAR_EPSILON)
+    {
+        d_cosTheta += sum(d_out * (f90 - f0) * -5.0f * powf(1.0f - cosTheta, 4.0f));
+    }
+}
+
+//------------------------------------------------------------------------
+// Ndf GGX
+
+__device__ float fwdNdfGGX(const float alphaSqr, const float cosTheta)
+{
+    float _cosTheta = clamp(cosTheta, SPECULAR_EPSILON, 1.0f - SPECULAR_EPSILON);
+    float d = (_cosTheta * alphaSqr - _cosTheta) * _cosTheta + 1.0f;
+    return alphaSqr / (d * d * M_PI);
+}
+
+__device__ void bwdNdfGGX(const float alphaSqr, const float cosTheta, float& d_alphaSqr, float& d_cosTheta, const float d_out)
+{
+    // Torch only back propagates if clamp doesn't trigger
+    float _cosTheta = clamp(cosTheta, SPECULAR_EPSILON, 1.0f - SPECULAR_EPSILON);
+    float cosThetaSqr = _cosTheta * _cosTheta;
+    d_alphaSqr += d_out * (1.0f - (alphaSqr + 1.0f) * cosThetaSqr) / (M_PI * powf((alphaSqr - 1.0) * cosThetaSqr + 1.0f, 3.0f));
+    if (cosTheta > SPECULAR_EPSILON && cosTheta < 1.0f - SPECULAR_EPSILON)
+    {
+        d_cosTheta += d_out * -(4.0f * (alphaSqr - 1.0f) * alphaSqr * cosTheta) / (M_PI * powf((alphaSqr - 1.0) * cosThetaSqr + 1.0f, 3.0f));
+    }
+}
+
+//------------------------------------------------------------------------
+// Lambda GGX
+
+__device__ float fwdLambdaGGX(const float alphaSqr, const float cosTheta)
+{
+    float _cosTheta = clamp(cosTheta, SPECULAR_EPSILON, 1.0f - SPECULAR_EPSILON);
+    float cosThetaSqr = _cosTheta * _cosTheta;
+    float tanThetaSqr = (1.0 - cosThetaSqr) / cosThetaSqr;
+    float res = 0.5f * (sqrtf(1.0f + alphaSqr * tanThetaSqr) - 1.0f);
+    return res;
+}
+
+__device__ void bwdLambdaGGX(const float alphaSqr, const float cosTheta, float& d_alphaSqr, float& d_cosTheta, const float d_out)
+{
+    float _cosTheta = clamp(cosTheta, SPECULAR_EPSILON, 1.0f - SPECULAR_EPSILON);
+    float cosThetaSqr = _cosTheta * _cosTheta;
+    float tanThetaSqr = (1.0 - cosThetaSqr) / cosThetaSqr;
+    float res = 0.5f * (sqrtf(1.0f + alphaSqr * tanThetaSqr) - 1.0f);
+
+    d_alphaSqr += d_out * (0.25 * tanThetaSqr) / sqrtf(alphaSqr * tanThetaSqr + 1.0f);
+    if (cosTheta > SPECULAR_EPSILON && cosTheta < 1.0f - SPECULAR_EPSILON)
+        d_cosTheta += d_out * -(0.5 * alphaSqr) / (powf(_cosTheta, 3.0f) * sqrtf(alphaSqr / cosThetaSqr - alphaSqr + 1.0f));
+}
+
+//------------------------------------------------------------------------
+// Masking GGX
+
+__device__ float fwdMaskingSmithGGXCorrelated(const float alphaSqr, const float cosThetaI, const float cosThetaO)
+{
+    float lambdaI = fwdLambdaGGX(alphaSqr, cosThetaI);
+    float lambdaO = fwdLambdaGGX(alphaSqr, cosThetaO);
+    return 1.0f / (1.0f + lambdaI + lambdaO);
+}
+
+__device__ void bwdMaskingSmithGGXCorrelated(const float alphaSqr, const float cosThetaI, const float cosThetaO, float& d_alphaSqr, float& d_cosThetaI, float& d_cosThetaO, const float d_out)
+{
+    // FWD eval
+    float lambdaI = fwdLambdaGGX(alphaSqr, cosThetaI);
+    float lambdaO = fwdLambdaGGX(alphaSqr, cosThetaO);
+
+    // BWD eval
+    float d_lambdaIO = -d_out / powf(1.0f + lambdaI + lambdaO, 2.0f);
+    bwdLambdaGGX(alphaSqr, cosThetaI, d_alphaSqr, d_cosThetaI, d_lambdaIO);
+    bwdLambdaGGX(alphaSqr, cosThetaO, d_alphaSqr, d_cosThetaO, d_lambdaIO);
+}
+
+//------------------------------------------------------------------------
+// GGX specular
+
+__device__ vec3f fwdPbrSpecular(const vec3f col, const vec3f nrm, const vec3f wo, const vec3f wi, const float alpha, const float min_roughness)
+{
+    float _alpha = clamp(alpha, min_roughness * min_roughness, 1.0f);
+    float alphaSqr = _alpha * _alpha;
+
+    vec3f h = safeNormalize(wo + wi);
+    float woDotN = dot(wo, nrm);
+    float wiDotN = dot(wi, nrm);
+    float woDotH = dot(wo, h);
+    float nDotH = dot(nrm, h);
+
+    float D = fwdNdfGGX(alphaSqr, nDotH);
+    float G = fwdMaskingSmithGGXCorrelated(alphaSqr, woDotN, wiDotN);
+    vec3f F = fwdFresnelSchlick(col, 1.0f, woDotH);
+    vec3f w = F * D * G * 0.25 / woDotN;
+
+    bool frontfacing = (woDotN > SPECULAR_EPSILON) & (wiDotN > SPECULAR_EPSILON);
+    return frontfacing ? w : 0.0f;
+}
+
+__device__ void bwdPbrSpecular(
+    const vec3f col, const vec3f nrm, const vec3f wo, const vec3f wi, const float alpha, const float min_roughness,
+    vec3f& d_col, vec3f& d_nrm, vec3f& d_wo, vec3f& d_wi, float& d_alpha, const vec3f d_out)
+{
+    ///////////////////////////////////////////////////////////////////////
+    // FWD eval
+
+    float _alpha = clamp(alpha, min_roughness * min_roughness, 1.0f);
+    float alphaSqr = _alpha * _alpha;
+
+    vec3f h = safeNormalize(wo + wi);
+    float woDotN = dot(wo, nrm);
+    float wiDotN = dot(wi, nrm);
+    float woDotH = dot(wo, h);
+    float nDotH = dot(nrm, h);
+
+    float D = fwdNdfGGX(alphaSqr, nDotH);
+    float G = fwdMaskingSmithGGXCorrelated(alphaSqr, woDotN, wiDotN);
+    vec3f F = fwdFresnelSchlick(col, 1.0f, woDotH);
+    vec3f w = F * D * G * 0.25 / woDotN;
+    bool frontfacing = (woDotN > SPECULAR_EPSILON) & (wiDotN > SPECULAR_EPSILON);
+
+    if (frontfacing)
+    {
+        ///////////////////////////////////////////////////////////////////////
+        // BWD eval
+
+        vec3f d_F = d_out * D * G * 0.25f / woDotN;
+        float d_D = sum(d_out * F * G * 0.25f / woDotN);
+        float d_G = sum(d_out * F * D * 0.25f / woDotN);
+
+        float d_woDotN = -sum(d_out * F * D * G * 0.25f / (woDotN * woDotN));
+
+        vec3f d_f90(0);
+        float d_woDotH(0), d_wiDotN(0), d_nDotH(0), d_alphaSqr(0);
+        bwdFresnelSchlick(col, 1.0f, woDotH, d_col, d_f90, d_woDotH, d_F);
+        bwdMaskingSmithGGXCorrelated(alphaSqr, woDotN, wiDotN, d_alphaSqr, d_woDotN, d_wiDotN, d_G);
+        bwdNdfGGX(alphaSqr, nDotH, d_alphaSqr, d_nDotH, d_D);
+
+        vec3f d_h(0);
+        bwdDot(nrm, h, d_nrm, d_h, d_nDotH);
+        bwdDot(wo, h, d_wo, d_h, d_woDotH);
+        bwdDot(wi, nrm, d_wi, d_nrm, d_wiDotN);
+        bwdDot(wo, nrm, d_wo, d_nrm, d_woDotN);
+
+        vec3f d_h_unnorm(0);
+        bwdSafeNormalize(wo + wi, d_h_unnorm, d_h);
+        d_wo += d_h_unnorm;
+        d_wi += d_h_unnorm;
+
+        if (alpha > min_roughness * min_roughness)
+            d_alpha += d_alphaSqr * 2 * alpha;
+    }
+}
+
+//------------------------------------------------------------------------
+// Full PBR BSDF
+
+__device__ vec3f fwdPbrBSDF(const vec3f kd, const vec3f arm, const vec3f pos, const vec3f nrm, const vec3f view_pos, const vec3f light_pos, const float min_roughness)
+{
+    vec3f wo = safeNormalize(view_pos - pos);
+    vec3f wi = safeNormalize(light_pos - pos);
+
+    float alpha = arm.y * arm.y;
+    vec3f spec_col = (0.04f * (1.0f - arm.z) + kd * arm.z) * (1.0 - arm.x);
+    vec3f diff_col = kd * (1.0f - arm.z);
+
+    float lambert = fwdLambert(nrm, wi);
+    vec3f diffuse = diff_col * lambert;
+    vec3f specular = fwdPbrSpecular(spec_col, nrm, wo, wi, alpha, min_roughness);
+
+    return diffuse + specular;
+}
+
+__device__ void bwdPbrBSDF(
+    const vec3f kd, const vec3f arm, const vec3f pos, const vec3f nrm, const vec3f view_pos, const vec3f light_pos, const float min_roughness,
+    vec3f& d_kd, vec3f& d_arm, vec3f& d_pos, vec3f& d_nrm, vec3f& d_view_pos, vec3f& d_light_pos, const vec3f d_out)
+{
+    ////////////////////////////////////////////////////////////////////////
+    // FWD
+    vec3f _wi = light_pos - pos;
+    vec3f _wo = view_pos - pos;
+    vec3f wi = safeNormalize(_wi);
+    vec3f wo = safeNormalize(_wo);
+
+    float alpha = arm.y * arm.y;
+    vec3f spec_col = (0.04f * (1.0f - arm.z) + kd * arm.z) * (1.0 - arm.x);
+    vec3f diff_col = kd * (1.0f - arm.z);
+    float lambert = fwdLambert(nrm, wi);
+
+    ////////////////////////////////////////////////////////////////////////
+    // BWD
+
+    float d_alpha(0);
+    vec3f d_spec_col(0), d_wi(0), d_wo(0);
+    bwdPbrSpecular(spec_col, nrm, wo, wi, alpha, min_roughness, d_spec_col, d_nrm, d_wo, d_wi, d_alpha, d_out);
+
+    float d_lambert = sum(diff_col * d_out);
+    bwdLambert(nrm, wi, d_nrm, d_wi, d_lambert);
+
+    // Backprop: diff_col = kd * (1.0f - arm.z)
+    vec3f d_diff_col = d_out * lambert;
+    d_kd += d_diff_col * (1.0f - arm.z);
+    d_arm.z -= sum(d_diff_col * kd);
+
+    // Backprop: spec_col = (0.04f * (1.0f - arm.z) + kd * arm.z) * (1.0 - arm.x)
+    d_kd -= d_spec_col * (arm.x - 1.0f) * arm.z;
+    d_arm.x += sum(d_spec_col * (arm.z * (0.04f - kd) - 0.04f));
+    d_arm.z -= sum(d_spec_col * (kd - 0.04f) * (arm.x - 1.0f));
+
+    // Backprop: alpha = arm.y * arm.y
+    d_arm.y += d_alpha * 2 * arm.y;
+
+    // Backprop: vec3f wi = safeNormalize(light_pos - pos);
+    vec3f d__wi(0);
+    bwdSafeNormalize(_wi, d__wi, d_wi);
+    d_light_pos += d__wi;
+    d_pos -= d__wi;
+
+    // Backprop: vec3f wo = safeNormalize(view_pos - pos);
+    vec3f d__wo(0);
+    bwdSafeNormalize(_wo, d__wo, d_wo);
+    d_view_pos += d__wo;
+    d_pos -= d__wo;
+}
+
+//------------------------------------------------------------------------
+// Kernels
+
+__global__ void LambertFwdKernel(LambertKernelParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f nrm = p.nrm.fetch3(px, py, pz);
+    vec3f wi = p.wi.fetch3(px, py, pz);
+
+    float res = fwdLambert(nrm, wi);
+
+    p.out.store(px, py, pz, res);
+}
+
+__global__ void LambertBwdKernel(LambertKernelParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f nrm = p.nrm.fetch3(px, py, pz);
+    vec3f wi = p.wi.fetch3(px, py, pz);
+    float d_out = p.out.fetch1(px, py, pz);
+
+    vec3f d_nrm(0), d_wi(0);
+    bwdLambert(nrm, wi, d_nrm, d_wi, d_out);
+
+    p.nrm.store_grad(px, py, pz, d_nrm);
+    p.wi.store_grad(px, py, pz, d_wi);
+}
+
+__global__ void FresnelShlickFwdKernel(FresnelShlickKernelParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f f0 = p.f0.fetch3(px, py, pz);
+    vec3f f90 = p.f90.fetch3(px, py, pz);
+    float cosTheta = p.cosTheta.fetch1(px, py, pz);
+
+    vec3f res = fwdFresnelSchlick(f0, f90, cosTheta);
+    p.out.store(px, py, pz, res);
+}
+
+__global__ void FresnelShlickBwdKernel(FresnelShlickKernelParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f f0 = p.f0.fetch3(px, py, pz);
+    vec3f f90 = p.f90.fetch3(px, py, pz);
+    float cosTheta = p.cosTheta.fetch1(px, py, pz);
+    vec3f d_out = p.out.fetch3(px, py, pz);
+
+    vec3f d_f0(0), d_f90(0);
+    float d_cosTheta(0);
+    bwdFresnelSchlick(f0, f90, cosTheta, d_f0, d_f90, d_cosTheta, d_out);
+
+    p.f0.store_grad(px, py, pz, d_f0);
+    p.f90.store_grad(px, py, pz, d_f90);
+    p.cosTheta.store_grad(px, py, pz, d_cosTheta);
+}
+
+__global__ void ndfGGXFwdKernel(NdfGGXParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    float alphaSqr = p.alphaSqr.fetch1(px, py, pz);
+    float cosTheta = p.cosTheta.fetch1(px, py, pz);
+    float res = fwdNdfGGX(alphaSqr, cosTheta);
+    
+    p.out.store(px, py, pz, res);
+}
+
+__global__ void ndfGGXBwdKernel(NdfGGXParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    float alphaSqr = p.alphaSqr.fetch1(px, py, pz);
+    float cosTheta = p.cosTheta.fetch1(px, py, pz);
+    float d_out = p.out.fetch1(px, py, pz);
+
+    float d_alphaSqr(0), d_cosTheta(0);
+    bwdNdfGGX(alphaSqr, cosTheta, d_alphaSqr, d_cosTheta, d_out);
+
+    p.alphaSqr.store_grad(px, py, pz, d_alphaSqr);
+    p.cosTheta.store_grad(px, py, pz, d_cosTheta);
+}
+
+__global__ void lambdaGGXFwdKernel(NdfGGXParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    float alphaSqr = p.alphaSqr.fetch1(px, py, pz);
+    float cosTheta = p.cosTheta.fetch1(px, py, pz);
+    float res = fwdLambdaGGX(alphaSqr, cosTheta);
+
+    p.out.store(px, py, pz, res);
+}
+
+__global__ void lambdaGGXBwdKernel(NdfGGXParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    float alphaSqr = p.alphaSqr.fetch1(px, py, pz);
+    float cosTheta = p.cosTheta.fetch1(px, py, pz);
+    float d_out = p.out.fetch1(px, py, pz);
+
+    float d_alphaSqr(0), d_cosTheta(0);
+    bwdLambdaGGX(alphaSqr, cosTheta, d_alphaSqr, d_cosTheta, d_out);
+
+    p.alphaSqr.store_grad(px, py, pz, d_alphaSqr);
+    p.cosTheta.store_grad(px, py, pz, d_cosTheta);
+}
+
+__global__ void maskingSmithFwdKernel(MaskingSmithParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    float alphaSqr = p.alphaSqr.fetch1(px, py, pz);
+    float cosThetaI = p.cosThetaI.fetch1(px, py, pz);
+    float cosThetaO = p.cosThetaO.fetch1(px, py, pz);
+    float res = fwdMaskingSmithGGXCorrelated(alphaSqr, cosThetaI, cosThetaO);
+    
+    p.out.store(px, py, pz, res);
+}
+
+__global__ void maskingSmithBwdKernel(MaskingSmithParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    float alphaSqr = p.alphaSqr.fetch1(px, py, pz);
+    float cosThetaI = p.cosThetaI.fetch1(px, py, pz);
+    float cosThetaO = p.cosThetaO.fetch1(px, py, pz);
+    float d_out = p.out.fetch1(px, py, pz);
+
+    float d_alphaSqr(0), d_cosThetaI(0), d_cosThetaO(0);
+    bwdMaskingSmithGGXCorrelated(alphaSqr, cosThetaI, cosThetaO, d_alphaSqr, d_cosThetaI, d_cosThetaO, d_out);
+
+    p.alphaSqr.store_grad(px, py, pz, d_alphaSqr);
+    p.cosThetaI.store_grad(px, py, pz, d_cosThetaI);
+    p.cosThetaO.store_grad(px, py, pz, d_cosThetaO);
+}
+
+__global__ void pbrSpecularFwdKernel(PbrSpecular p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f col = p.col.fetch3(px, py, pz);
+    vec3f nrm = p.nrm.fetch3(px, py, pz);
+    vec3f wo = p.wo.fetch3(px, py, pz);
+    vec3f wi = p.wi.fetch3(px, py, pz);
+    float alpha = p.alpha.fetch1(px, py, pz);
+
+    vec3f res = fwdPbrSpecular(col, nrm, wo, wi, alpha, p.min_roughness);
+
+    p.out.store(px, py, pz, res);
+}
+
+__global__ void pbrSpecularBwdKernel(PbrSpecular p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f col = p.col.fetch3(px, py, pz);
+    vec3f nrm = p.nrm.fetch3(px, py, pz);
+    vec3f wo = p.wo.fetch3(px, py, pz);
+    vec3f wi = p.wi.fetch3(px, py, pz);
+    float alpha = p.alpha.fetch1(px, py, pz);
+    vec3f d_out = p.out.fetch3(px, py, pz);
+
+    float d_alpha(0);
+    vec3f d_col(0), d_nrm(0), d_wo(0), d_wi(0);
+    bwdPbrSpecular(col, nrm, wo, wi, alpha, p.min_roughness, d_col, d_nrm, d_wo, d_wi, d_alpha, d_out);
+
+    p.col.store_grad(px, py, pz, d_col);
+    p.nrm.store_grad(px, py, pz, d_nrm);
+    p.wo.store_grad(px, py, pz, d_wo);
+    p.wi.store_grad(px, py, pz, d_wi);
+    p.alpha.store_grad(px, py, pz, d_alpha);
+}
+
+__global__ void pbrBSDFFwdKernel(PbrBSDF p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f kd = p.kd.fetch3(px, py, pz);
+    vec3f arm = p.arm.fetch3(px, py, pz);
+    vec3f pos = p.pos.fetch3(px, py, pz);
+    vec3f nrm = p.nrm.fetch3(px, py, pz);
+    vec3f view_pos = p.view_pos.fetch3(px, py, pz);
+    vec3f light_pos = p.light_pos.fetch3(px, py, pz);
+
+    vec3f res = fwdPbrBSDF(kd, arm, pos, nrm, view_pos, light_pos, p.min_roughness);
+
+    p.out.store(px, py, pz, res);
+}
+__global__ void pbrBSDFBwdKernel(PbrBSDF p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f kd = p.kd.fetch3(px, py, pz);
+    vec3f arm = p.arm.fetch3(px, py, pz);
+    vec3f pos = p.pos.fetch3(px, py, pz);
+    vec3f nrm = p.nrm.fetch3(px, py, pz);
+    vec3f view_pos = p.view_pos.fetch3(px, py, pz);
+    vec3f light_pos = p.light_pos.fetch3(px, py, pz);
+    vec3f d_out = p.out.fetch3(px, py, pz);
+
+    vec3f d_kd(0), d_arm(0), d_pos(0), d_nrm(0), d_view_pos(0), d_light_pos(0);
+    bwdPbrBSDF(kd, arm, pos, nrm, view_pos, light_pos, p.min_roughness, d_kd, d_arm, d_pos, d_nrm, d_view_pos, d_light_pos, d_out);
+
+    p.kd.store_grad(px, py, pz, d_kd);
+    p.arm.store_grad(px, py, pz, d_arm);
+    p.pos.store_grad(px, py, pz, d_pos);
+    p.nrm.store_grad(px, py, pz, d_nrm);
+    p.view_pos.store_grad(px, py, pz, d_view_pos);
+    p.light_pos.store_grad(px, py, pz, d_light_pos);
+}

nvdiffmodeling/src/renderutils/c_src/bsdf.h

@@ -0,0 +1,70 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+
+#include "common.h"
+
+struct LambertKernelParams
+{
+    Tensor  nrm;
+    Tensor  wi;
+    Tensor  out;
+    dim3    gridSize;
+};
+
+struct FresnelShlickKernelParams
+{
+    Tensor  f0;
+    Tensor  f90;
+    Tensor  cosTheta;
+    Tensor  out;
+    dim3    gridSize;
+};
+
+struct NdfGGXParams
+{
+    Tensor  alphaSqr;
+    Tensor  cosTheta;
+    Tensor  out;
+    dim3    gridSize;
+};
+
+struct MaskingSmithParams
+{
+    Tensor  alphaSqr;
+    Tensor  cosThetaI;
+    Tensor  cosThetaO;
+    Tensor  out;
+    dim3    gridSize;
+};
+
+struct PbrSpecular
+{
+    Tensor  col;
+    Tensor  nrm;
+    Tensor  wo;
+    Tensor  wi;
+    Tensor  alpha;
+    Tensor  out;
+    dim3    gridSize;
+    float   min_roughness;
+};
+
+struct PbrBSDF
+{
+    Tensor  kd;
+    Tensor  arm;
+    Tensor  pos;
+    Tensor  nrm;
+    Tensor  view_pos;
+    Tensor  light_pos;
+    Tensor  out;
+    dim3    gridSize;
+    float   min_roughness;
+};

nvdiffmodeling/src/renderutils/c_src/common.cpp

@@ -0,0 +1,71 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <cuda_runtime.h>
+#include <algorithm>
+
+//------------------------------------------------------------------------
+// Block and grid size calculators for kernel launches.
+
+dim3 getLaunchBlockSize(int maxWidth, int maxHeight, dim3 dims)
+{
+    int maxThreads = maxWidth * maxHeight;
+    if (maxThreads <= 1 || (dims.x * dims.y) <= 1)
+        return dim3(1, 1, 1); // Degenerate.
+
+    // Start from max size.
+    int bw = maxWidth;
+    int bh = maxHeight;
+
+    // Optimizations for weirdly sized buffers.
+    if (dims.x < bw)
+    {
+        // Decrease block width to smallest power of two that covers the buffer width.
+        while ((bw >> 1) >= dims.x)
+            bw >>= 1;
+
+        // Maximize height.
+        bh = maxThreads / bw;
+        if (bh > dims.y)
+            bh = dims.y;
+    }
+    else if (dims.y < bh)
+    {
+        // Halve height and double width until fits completely inside buffer vertically.
+        while (bh > dims.y)
+        {
+            bh >>= 1;
+            if (bw < dims.x)
+                bw <<= 1;
+        }
+    }
+
+    // Done.
+    return dim3(bw, bh, 1);
+}
+
+// returns the size of a block that can be reduced using horizontal SIMD operations (e.g. __shfl_xor_sync)
+dim3 getWarpSize(dim3 blockSize)
+{
+    return dim3(
+        std::min(blockSize.x, 32u), 
+        std::min(std::max(32u / blockSize.x, 1u), std::min(32u, blockSize.y)), 
+        std::min(std::max(32u / (blockSize.x * blockSize.y), 1u), std::min(32u, blockSize.z))
+    );
+}
+
+dim3 getLaunchGridSize(dim3 blockSize, dim3 dims)
+{
+    dim3 gridSize;
+    gridSize.x = (dims.x  - 1) / blockSize.x + 1;
+    gridSize.y = (dims.y - 1) / blockSize.y + 1;
+    gridSize.z = (dims.z  - 1) / blockSize.z + 1;
+    return gridSize;
+}
+
+//------------------------------------------------------------------------

nvdiffmodeling/src/renderutils/c_src/common.h

@@ -0,0 +1,38 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+#include <cuda.h>
+#include <stdint.h>
+
+#include "vec3f.h"
+#include "vec4f.h"
+#include "tensor.h"
+
+dim3 getLaunchBlockSize(int maxWidth, int maxHeight, dim3 dims);
+dim3 getLaunchGridSize(dim3 blockSize, dim3 dims);
+
+#ifdef __CUDACC__
+
+#ifdef _MSC_VER
+#define M_PI 3.14159265358979323846f
+#endif
+
+__host__ __device__ static inline dim3 getWarpSize(dim3 blockSize)
+{
+    return dim3(
+        min(blockSize.x, 32u),
+        min(max(32u / blockSize.x, 1u), min(32u, blockSize.y)),
+        min(max(32u / (blockSize.x * blockSize.y), 1u), min(32u, blockSize.z))
+    );
+}
+
+__device__ static inline float clamp(float val, float mn, float mx) { return min(max(val, mn), mx); }
+#else
+dim3 getWarpSize(dim3 blockSize);
+#endif

nvdiffmodeling/src/renderutils/c_src/loss.cu

@@ -0,0 +1,207 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <cuda.h>
+
+#include "common.h"
+#include "loss.h"
+
+//------------------------------------------------------------------------
+// Utils
+
+__device__ inline float bwdAbs(float x) { return x == 0.0f ? 0.0f : x < 0.0f ? -1.0f : 1.0f; }
+
+__device__ float warpSum(float val) {
+    for (int i = 1; i < 32; i *= 2)
+        val += __shfl_xor_sync(0xFFFFFFFF, val, i);
+    return val;
+}
+
+//------------------------------------------------------------------------
+// Tonemapping
+
+__device__ inline float fwdSRGB(float x)
+{
+    return x > 0.0031308f ? powf(max(x, 0.0031308f), 1.0f / 2.4f) * 1.055f - 0.055f : 12.92f * max(x, 0.0f);
+}
+
+__device__ inline void bwdSRGB(float x, float &d_x, float d_out)
+{
+    if (x > 0.0031308f)
+        d_x += d_out * 0.439583f / powf(x, 0.583333f);
+    else if (x > 0.0f)
+        d_x += d_out * 12.92f;
+}
+
+__device__ inline vec3f fwdTonemapLogSRGB(vec3f x)
+{
+    return vec3f(fwdSRGB(logf(x.x + 1.0f)), fwdSRGB(logf(x.y + 1.0f)), fwdSRGB(logf(x.z + 1.0f)));
+}
+
+__device__ inline void bwdTonemapLogSRGB(vec3f x, vec3f& d_x, vec3f d_out)
+{
+    if (x.x > 0.0f && x.x < 65535.0f)
+    {
+        bwdSRGB(logf(x.x + 1.0f), d_x.x, d_out.x);
+        d_x.x *= 1 / (x.x + 1.0f);
+    }
+    if (x.y > 0.0f && x.y < 65535.0f)
+    {
+        bwdSRGB(logf(x.y + 1.0f), d_x.y, d_out.y);
+        d_x.y *= 1 / (x.y + 1.0f);
+    }
+    if (x.z > 0.0f && x.z < 65535.0f)
+    {
+        bwdSRGB(logf(x.z + 1.0f), d_x.z, d_out.z);
+        d_x.z *= 1 / (x.z + 1.0f);
+    }
+}
+
+__device__ inline float fwdRELMSE(float img, float target, float eps = 0.1f)
+{
+    return (img - target) * (img - target) / (img * img + target * target + eps);
+}
+
+__device__ inline void bwdRELMSE(float img, float target, float &d_img, float &d_target, float d_out, float eps = 0.1f)
+{
+    float denom  = (target * target + img * img + eps);
+    d_img    += d_out * 2 * (img - target) * (target * (target + img) + eps) / (denom * denom);
+    d_target -= d_out * 2 * (img - target) * (img * (target + img) + eps) / (denom * denom);
+}
+
+__device__ inline float fwdSMAPE(float img, float target, float eps=0.01f)
+{
+    return abs(img - target) / (img + target + eps);
+}
+
+__device__ inline void bwdSMAPE(float img, float target, float& d_img, float& d_target, float d_out, float eps = 0.01f)
+{
+    float denom = (target + img + eps);
+    d_img    += d_out * bwdAbs(img - target) * (2 * target + eps) / (denom * denom);
+    d_target -= d_out * bwdAbs(img - target) * (2 * img + eps) / (denom * denom);
+}
+
+//------------------------------------------------------------------------
+// Kernels
+
+__global__ void imgLossFwdKernel(LossKernelParams p)
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+
+    float floss = 0.0f;
+    if (px < p.gridSize.x && py < p.gridSize.y && pz < p.gridSize.z)
+    {
+        vec3f img = p.img.fetch3(px, py, pz);
+        vec3f target = p.target.fetch3(px, py, pz);
+
+        img = vec3f(clamp(img.x, 0.0f, 65535.0f), clamp(img.y, 0.0f, 65535.0f), clamp(img.z, 0.0f, 65535.0f));
+        target = vec3f(clamp(target.x, 0.0f, 65535.0f), clamp(target.y, 0.0f, 65535.0f), clamp(target.z, 0.0f, 65535.0f));
+
+        if (p.tonemapper == TONEMAPPER_LOG_SRGB)
+        {
+            img = fwdTonemapLogSRGB(img);
+            target = fwdTonemapLogSRGB(target);
+        }
+
+        vec3f vloss(0);
+        if (p.loss == LOSS_MSE)
+            vloss = (img - target) * (img - target);
+        else if (p.loss == LOSS_RELMSE)
+            vloss = vec3f(fwdRELMSE(img.x, target.x), fwdRELMSE(img.y, target.y), fwdRELMSE(img.z, target.z));
+        else if (p.loss == LOSS_SMAPE)
+            vloss = vec3f(fwdSMAPE(img.x, target.x), fwdSMAPE(img.y, target.y), fwdSMAPE(img.z, target.z));
+        else
+            vloss = vec3f(abs(img.x - target.x), abs(img.y - target.y), abs(img.z - target.z));
+        
+        floss = sum(vloss) / 3.0f;
+    }
+
+    floss = warpSum(floss);
+
+    dim3 warpSize = getWarpSize(blockDim);
+    if (px < p.gridSize.x && py < p.gridSize.y && pz < p.gridSize.z && threadIdx.x % warpSize.x == 0 && threadIdx.y % warpSize.y == 0 && threadIdx.z % warpSize.z == 0)
+        p.out.store(px / warpSize.x, py / warpSize.y, pz / warpSize.z, floss);
+}
+
+__global__ void imgLossBwdKernel(LossKernelParams p)
+{ 
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    dim3 warpSize = getWarpSize(blockDim);
+
+    vec3f _img = p.img.fetch3(px, py, pz);
+    vec3f _target = p.target.fetch3(px, py, pz);
+    float d_out = p.out.fetch1(px / warpSize.x, py / warpSize.y, pz / warpSize.z);
+
+    /////////////////////////////////////////////////////////////////////
+    // FWD
+
+    vec3f img = _img, target = _target;
+    if (p.tonemapper == TONEMAPPER_LOG_SRGB)
+    {
+        img = fwdTonemapLogSRGB(img);
+        target = fwdTonemapLogSRGB(target);
+    }
+
+    /////////////////////////////////////////////////////////////////////
+    // BWD
+
+    vec3f d_vloss = vec3f(d_out, d_out, d_out) / 3.0f;
+
+    vec3f d_img(0), d_target(0);
+    if (p.loss == LOSS_MSE)
+    {
+        d_img = vec3f(d_vloss.x * 2 * (img.x - target.x), d_vloss.y * 2 * (img.y - target.y), d_vloss.x * 2 * (img.z - target.z));
+        d_target = -d_img;
+    }
+    else if (p.loss == LOSS_RELMSE)
+    {
+        bwdRELMSE(img.x, target.x, d_img.x, d_target.x, d_vloss.x);
+        bwdRELMSE(img.y, target.y, d_img.y, d_target.y, d_vloss.y);
+        bwdRELMSE(img.z, target.z, d_img.z, d_target.z, d_vloss.z);
+    }
+    else if (p.loss == LOSS_SMAPE)
+    {
+        bwdSMAPE(img.x, target.x, d_img.x, d_target.x, d_vloss.x);
+        bwdSMAPE(img.y, target.y, d_img.y, d_target.y, d_vloss.y);
+        bwdSMAPE(img.z, target.z, d_img.z, d_target.z, d_vloss.z);
+    }
+    else
+    {
+        d_img = d_vloss * vec3f(bwdAbs(img.x - target.x), bwdAbs(img.y - target.y), bwdAbs(img.z - target.z));
+        d_target = -d_img;
+    }
+
+
+    if (p.tonemapper == TONEMAPPER_LOG_SRGB)
+    {
+        vec3f d__img(0), d__target(0);
+        bwdTonemapLogSRGB(_img, d__img, d_img);
+        bwdTonemapLogSRGB(_target, d__target, d_target);
+        d_img = d__img; d_target = d__target;
+    }
+
+    if (_img.x <= 0.0f || _img.x >= 65535.0f) d_img.x = 0;
+    if (_img.y <= 0.0f || _img.y >= 65535.0f) d_img.y = 0;
+    if (_img.z <= 0.0f || _img.z >= 65535.0f) d_img.z = 0;
+    if (_target.x <= 0.0f || _target.x >= 65535.0f) d_target.x = 0;
+    if (_target.y <= 0.0f || _target.y >= 65535.0f) d_target.y = 0;
+    if (_target.z <= 0.0f || _target.z >= 65535.0f) d_target.z = 0;
+
+    p.img.store_grad(px, py, pz, d_img);
+    p.target.store_grad(px, py, pz, d_target);
+}

nvdiffmodeling/src/renderutils/c_src/loss.h

@@ -0,0 +1,35 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+
+#include "common.h"
+
+enum TonemapperType
+{
+    TONEMAPPER_NONE = 0,
+    TONEMAPPER_LOG_SRGB = 1
+};
+
+enum LossType
+{
+    LOSS_L1 = 0,
+    LOSS_MSE = 1,
+    LOSS_RELMSE = 2,
+    LOSS_SMAPE = 3
+};
+
+struct LossKernelParams
+{
+    Tensor          img;
+    Tensor          target;
+    Tensor          out;
+    dim3            gridSize;
+    TonemapperType  tonemapper;
+    LossType        loss;
+};

nvdiffmodeling/src/renderutils/c_src/mesh.cu

@@ -0,0 +1,90 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <cuda.h>
+
+#include "common.h"
+#include "mesh.h"
+
+
+//------------------------------------------------------------------------
+// Kernels
+
+__global__ void xfmPointsFwdKernel(XfmKernelParams p)
+{
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int pz = blockIdx.z * blockDim.z + threadIdx.z;
+
+    __shared__ float mtx[4][4];
+    if (threadIdx.x < 16)
+        mtx[threadIdx.x % 4][threadIdx.x / 4] = p.matrix.fetch(p.matrix.nhwcIndex(pz, threadIdx.x / 4, threadIdx.x % 4, 0));
+    __syncthreads();
+    
+    if (px >= p.gridSize.x)
+        return;
+
+    vec3f pos(
+        p.points.fetch(p.points.nhwcIndex(pz, px, 0, 0)),
+        p.points.fetch(p.points.nhwcIndex(pz, px, 1, 0)),
+        p.points.fetch(p.points.nhwcIndex(pz, px, 2, 0))
+    );
+
+    if (p.isPoints)
+    {
+        p.out.store(p.out.nhwcIndex(pz, px, 0, 0), pos.x * mtx[0][0] + pos.y * mtx[1][0] + pos.z * mtx[2][0] + mtx[3][0]);
+        p.out.store(p.out.nhwcIndex(pz, px, 1, 0), pos.x * mtx[0][1] + pos.y * mtx[1][1] + pos.z * mtx[2][1] + mtx[3][1]);
+        p.out.store(p.out.nhwcIndex(pz, px, 2, 0), pos.x * mtx[0][2] + pos.y * mtx[1][2] + pos.z * mtx[2][2] + mtx[3][2]);
+        p.out.store(p.out.nhwcIndex(pz, px, 3, 0), pos.x * mtx[0][3] + pos.y * mtx[1][3] + pos.z * mtx[2][3] + mtx[3][3]);
+    }
+    else
+    {
+        p.out.store(p.out.nhwcIndex(pz, px, 0, 0), pos.x * mtx[0][0] + pos.y * mtx[1][0] + pos.z * mtx[2][0]);
+        p.out.store(p.out.nhwcIndex(pz, px, 1, 0), pos.x * mtx[0][1] + pos.y * mtx[1][1] + pos.z * mtx[2][1]);
+        p.out.store(p.out.nhwcIndex(pz, px, 2, 0), pos.x * mtx[0][2] + pos.y * mtx[1][2] + pos.z * mtx[2][2]);
+    }
+}
+
+__global__ void xfmPointsBwdKernel(XfmKernelParams p)
+{ 
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int pz = blockIdx.z * blockDim.z + threadIdx.z;
+
+    __shared__ float mtx[4][4];
+    if (threadIdx.x < 16)
+        mtx[threadIdx.x % 4][threadIdx.x / 4] = p.matrix.fetch(p.matrix.nhwcIndex(pz, threadIdx.x / 4, threadIdx.x % 4, 0));
+    __syncthreads();
+
+    if (px >= p.gridSize.x)
+        return;
+
+    vec3f pos(
+        p.points.fetch(p.points.nhwcIndex(pz, px, 0, 0)),
+        p.points.fetch(p.points.nhwcIndex(pz, px, 1, 0)),
+        p.points.fetch(p.points.nhwcIndex(pz, px, 2, 0))
+    );
+
+    vec4f d_out(
+        p.out.fetch(p.out.nhwcIndex(pz, px, 0, 0)),
+        p.out.fetch(p.out.nhwcIndex(pz, px, 1, 0)),
+        p.out.fetch(p.out.nhwcIndex(pz, px, 2, 0)),
+        p.out.fetch(p.out.nhwcIndex(pz, px, 3, 0))
+    );
+
+    if (p.isPoints)
+    {
+        p.points.store_grad(p.points._nhwcIndex(pz, px, 0, 0), d_out.x * mtx[0][0] + d_out.y * mtx[0][1] + d_out.z * mtx[0][2] + d_out.w * mtx[0][3]);
+        p.points.store_grad(p.points._nhwcIndex(pz, px, 1, 0), d_out.x * mtx[1][0] + d_out.y * mtx[1][1] + d_out.z * mtx[1][2] + d_out.w * mtx[1][3]);
+        p.points.store_grad(p.points._nhwcIndex(pz, px, 2, 0), d_out.x * mtx[2][0] + d_out.y * mtx[2][1] + d_out.z * mtx[2][2] + d_out.w * mtx[2][3]);
+    }
+    else
+    {
+        p.points.store_grad(p.points._nhwcIndex(pz, px, 0, 0), d_out.x * mtx[0][0] + d_out.y * mtx[0][1] + d_out.z * mtx[0][2]);
+        p.points.store_grad(p.points._nhwcIndex(pz, px, 1, 0), d_out.x * mtx[1][0] + d_out.y * mtx[1][1] + d_out.z * mtx[1][2]);
+        p.points.store_grad(p.points._nhwcIndex(pz, px, 2, 0), d_out.x * mtx[2][0] + d_out.y * mtx[2][1] + d_out.z * mtx[2][2]);
+    }
+}

nvdiffmodeling/src/renderutils/c_src/mesh.h

@@ -0,0 +1,20 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+
+#include "common.h"
+
+struct XfmKernelParams
+{
+    bool            isPoints;
+    Tensor          points;
+    Tensor          matrix;
+    Tensor          out;
+    dim3            gridSize;
+};

nvdiffmodeling/src/renderutils/c_src/normal.cu

@@ -0,0 +1,179 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include "common.h"
+#include "normal.h"
+
+#define NORMAL_THRESHOLD 0.1f
+
+//------------------------------------------------------------------------
+// Perturb shading normal by tangent frame
+
+__device__ vec3f fwdPerturbNormal(const vec3f perturbed_nrm, const vec3f smooth_nrm, const vec3f smooth_tng, bool opengl)
+{
+    vec3f _smooth_bitng = cross(smooth_tng, smooth_nrm);
+    vec3f smooth_bitng = safeNormalize(_smooth_bitng);
+    vec3f _shading_nrm = smooth_tng * perturbed_nrm.x + (opengl ? -1 : 1) * smooth_bitng * perturbed_nrm.y + smooth_nrm * max(perturbed_nrm.z, 0.0f);
+    return safeNormalize(_shading_nrm);
+}
+
+__device__ void bwdPerturbNormal(const vec3f perturbed_nrm, const vec3f smooth_nrm, const vec3f smooth_tng, vec3f &d_perturbed_nrm, vec3f &d_smooth_nrm, vec3f &d_smooth_tng, const vec3f d_out, bool opengl)
+{
+    ////////////////////////////////////////////////////////////////////////
+    // FWD
+    vec3f _smooth_bitng = cross(smooth_tng, smooth_nrm);
+    vec3f smooth_bitng = safeNormalize(_smooth_bitng);
+    vec3f _shading_nrm = smooth_tng * perturbed_nrm.x + (opengl ? -1 : 1) * smooth_bitng * perturbed_nrm.y + smooth_nrm * max(perturbed_nrm.z, 0.0f);
+        
+    ////////////////////////////////////////////////////////////////////////
+    // BWD
+    vec3f d_shading_nrm(0);
+    bwdSafeNormalize(_shading_nrm, d_shading_nrm, d_out);
+
+    vec3f d_smooth_bitng(0);
+    
+    if (perturbed_nrm.z > 0.0f)
+    {
+        d_smooth_nrm += d_shading_nrm * perturbed_nrm.z;
+        d_perturbed_nrm.z += sum(d_shading_nrm * smooth_nrm);
+    }
+
+    d_smooth_bitng += (opengl ? -1 : 1) * d_shading_nrm * perturbed_nrm.y;
+    d_perturbed_nrm.y += (opengl ? -1 : 1) * sum(d_shading_nrm * smooth_bitng);
+
+    d_smooth_tng += d_shading_nrm * perturbed_nrm.x;
+    d_perturbed_nrm.x += sum(d_shading_nrm * smooth_tng);
+
+    vec3f d__smooth_bitng(0);
+    bwdSafeNormalize(_smooth_bitng, d__smooth_bitng, d_smooth_bitng);
+
+    bwdCross(smooth_tng, smooth_nrm, d_smooth_tng, d_smooth_nrm, d__smooth_bitng);
+}
+
+//------------------------------------------------------------------------
+#define bent_nrm_eps 0.001f
+
+__device__ vec3f fwdBendNormal(const vec3f view_vec, const vec3f smooth_nrm, const vec3f geom_nrm)
+{
+    float dp = dot(view_vec, smooth_nrm);
+    float t = clamp(dp / NORMAL_THRESHOLD, 0.0f, 1.0f);
+    return geom_nrm * (1.0f - t) + smooth_nrm * t;
+}
+
+__device__ void bwdBendNormal(const vec3f view_vec, const vec3f smooth_nrm, const vec3f geom_nrm, vec3f& d_view_vec, vec3f& d_smooth_nrm, vec3f& d_geom_nrm, const vec3f d_out)
+{
+    ////////////////////////////////////////////////////////////////////////
+    // FWD
+    float dp = dot(view_vec, smooth_nrm);
+    float t = clamp(dp / NORMAL_THRESHOLD, 0.0f, 1.0f);
+
+    ////////////////////////////////////////////////////////////////////////
+    // BWD
+    if (dp > NORMAL_THRESHOLD)
+        d_smooth_nrm += d_out;
+    else
+    {
+        // geom_nrm * (1.0f - t) + smooth_nrm * t;
+        d_geom_nrm   += d_out * (1.0f - t);
+        d_smooth_nrm += d_out * t;
+        float d_t = sum(d_out * (smooth_nrm - geom_nrm));
+
+        float d_dp = dp < 0.0f || dp > NORMAL_THRESHOLD ? 0.0f : d_t / NORMAL_THRESHOLD;
+
+        bwdDot(view_vec, smooth_nrm, d_view_vec, d_smooth_nrm, d_dp);
+    }
+}
+
+//------------------------------------------------------------------------
+// Kernels
+
+__global__ void PrepareShadingNormalFwdKernel(PrepareShadingNormalKernelParams p) 
+{
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f pos = p.pos.fetch3(px, py, pz);
+    vec3f view_pos = p.view_pos.fetch3(px, py, pz);
+    vec3f perturbed_nrm = p.perturbed_nrm.fetch3(px, py, pz);
+    vec3f _smooth_nrm = p.smooth_nrm.fetch3(px, py, pz);
+    vec3f _smooth_tng = p.smooth_tng.fetch3(px, py, pz);
+    vec3f geom_nrm = p.geom_nrm.fetch3(px, py, pz);
+
+    vec3f smooth_nrm = safeNormalize(_smooth_nrm);
+    vec3f smooth_tng = safeNormalize(_smooth_tng);
+    vec3f view_vec = safeNormalize(view_pos - pos);
+    vec3f shading_nrm = fwdPerturbNormal(perturbed_nrm, smooth_nrm, smooth_tng, p.opengl);
+
+    vec3f res;
+    if (p.two_sided_shading && dot(view_vec, geom_nrm) < 0.0f)
+        res = fwdBendNormal(view_vec, -shading_nrm, -geom_nrm);
+    else
+        res = fwdBendNormal(view_vec, shading_nrm, geom_nrm);
+
+    p.out.store(px, py, pz, res);
+}
+
+__global__ void PrepareShadingNormalBwdKernel(PrepareShadingNormalKernelParams p) 
+{ 
+    // Calculate pixel position.
+    unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
+    unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
+    unsigned int pz = blockIdx.z;
+    if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
+        return;
+
+    vec3f pos = p.pos.fetch3(px, py, pz);
+    vec3f view_pos = p.view_pos.fetch3(px, py, pz);
+    vec3f perturbed_nrm = p.perturbed_nrm.fetch3(px, py, pz);
+    vec3f _smooth_nrm = p.smooth_nrm.fetch3(px, py, pz);
+    vec3f _smooth_tng = p.smooth_tng.fetch3(px, py, pz);
+    vec3f geom_nrm = p.geom_nrm.fetch3(px, py, pz);
+    vec3f d_out = p.out.fetch3(px, py, pz);
+
+    ///////////////////////////////////////////////////////////////////////////////////////////////////
+    // FWD
+
+    vec3f smooth_nrm = safeNormalize(_smooth_nrm);
+    vec3f smooth_tng = safeNormalize(_smooth_tng);
+    vec3f _view_vec = view_pos - pos;
+    vec3f view_vec = safeNormalize(view_pos - pos);
+
+    vec3f shading_nrm = fwdPerturbNormal(perturbed_nrm, smooth_nrm, smooth_tng, p.opengl);
+
+    ///////////////////////////////////////////////////////////////////////////////////////////////////
+    // BWD
+
+    vec3f d_view_vec(0), d_shading_nrm(0), d_geom_nrm(0);
+    if (p.two_sided_shading && dot(view_vec, geom_nrm) < 0.0f)
+    {
+        bwdBendNormal(view_vec, -shading_nrm, -geom_nrm, d_view_vec, d_shading_nrm, d_geom_nrm, d_out);
+        d_shading_nrm = -d_shading_nrm;
+        d_geom_nrm = -d_geom_nrm;
+    }
+    else
+        bwdBendNormal(view_vec, shading_nrm, geom_nrm, d_view_vec, d_shading_nrm, d_geom_nrm, d_out);
+
+    vec3f d_perturbed_nrm(0), d_smooth_nrm(0), d_smooth_tng(0);
+    bwdPerturbNormal(perturbed_nrm, smooth_nrm, smooth_tng, d_perturbed_nrm, d_smooth_nrm, d_smooth_tng, d_shading_nrm, p.opengl);
+
+    vec3f d__view_vec(0), d__smooth_nrm(0), d__smooth_tng(0);
+    bwdSafeNormalize(_view_vec, d__view_vec, d_view_vec);
+    bwdSafeNormalize(_smooth_nrm, d__smooth_nrm, d_smooth_nrm);
+    bwdSafeNormalize(_smooth_tng, d__smooth_tng, d_smooth_tng);
+
+    p.pos.store_grad(px, py, pz, -d__view_vec);
+    p.view_pos.store_grad(px, py, pz, d__view_vec);
+    p.perturbed_nrm.store_grad(px, py, pz, d_perturbed_nrm);
+    p.smooth_nrm.store_grad(px, py, pz, d__smooth_nrm);
+    p.smooth_tng.store_grad(px, py, pz, d__smooth_tng);
+    p.geom_nrm.store_grad(px, py, pz, d_geom_nrm);
+}

nvdiffmodeling/src/renderutils/c_src/normal.h

@@ -0,0 +1,24 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+
+#include "common.h"
+
+struct PrepareShadingNormalKernelParams
+{
+    Tensor  pos;
+    Tensor  view_pos;
+    Tensor  perturbed_nrm;
+    Tensor  smooth_nrm;
+    Tensor  smooth_tng;
+    Tensor  geom_nrm;
+    Tensor  out;
+    dim3    gridSize;
+    bool    two_sided_shading, opengl;
+};

nvdiffmodeling/src/renderutils/c_src/tensor.h

@@ -0,0 +1,86 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+#if defined(__CUDACC__) && defined(BFLOAT16)
+#include <cuda_bf16.h> // bfloat16 is float32 compatible with less mantissa bits
+#endif
+
+//---------------------------------------------------------------------------------
+// CUDA-side Tensor class for in/out parameter parsing. Can be float32 or bfloat16
+
+struct Tensor
+{
+    void*   val;
+    void*   d_val;
+    int     dims[4];
+    int     strides[4];
+    bool    fp16;
+    Tensor() : val(nullptr), d_val(nullptr), fp16(true), dims{ 0, 0, 0, 0 }, strides{ 0, 0, 0, 0 } {}
+
+#ifdef __CUDACC__
+    // Helpers to index and read/write a single element
+    __device__ inline int   _nhwcIndex(int n, int h, int w, int c) const { return n * strides[0] + h * strides[1] + w * strides[2] + c * strides[3]; }
+    __device__ inline int   nhwcIndex(int n, int h, int w, int c) const { return (dims[0] == 1 ? 0 : n * strides[0]) + (dims[1] == 1 ? 0 : h * strides[1]) + (dims[2] == 1 ? 0 : w * strides[2]) + (dims[3] == 1 ? 0 : c * strides[3]); }
+    __device__ inline int   nhwcIndexContinuous(int n, int h, int w, int c) const { return ((n * dims[1] + h) * dims[2] + w) * dims[3] + c; }
+#ifdef BFLOAT16
+    __device__ inline float fetch(unsigned int idx) const { return fp16 ? __bfloat162float(((__nv_bfloat16*)val)[idx]) : ((float*)val)[idx]; }
+    __device__ inline void  store(unsigned int idx, float _val) { if (fp16) ((__nv_bfloat16*)val)[idx] = __float2bfloat16(_val); else ((float*)val)[idx] = _val; }
+    __device__ inline void  store_grad(unsigned int idx, float _val) { if (fp16) ((__nv_bfloat16*)d_val)[idx] = __float2bfloat16(_val); else ((float*)d_val)[idx] = _val; }
+#else
+    __device__ inline float fetch(unsigned int idx) const { return ((float*)val)[idx]; }
+    __device__ inline void  store(unsigned int idx, float _val) { ((float*)val)[idx] = _val; }
+    __device__ inline void  store_grad(unsigned int idx, float _val) { ((float*)d_val)[idx] = _val; }
+#endif
+
+    //////////////////////////////////////////////////////////////////////////////////////////
+    // Fetch, use broadcasting for tensor dimensions of size 1
+    __device__ inline float fetch1(unsigned int x, unsigned int y, unsigned int z) const
+    {
+        return fetch(nhwcIndex(z, y, x, 0));
+    }
+
+    __device__ inline vec3f fetch3(unsigned int x, unsigned int y, unsigned int z) const
+    {
+        return vec3f(
+            fetch(nhwcIndex(z, y, x, 0)),
+            fetch(nhwcIndex(z, y, x, 1)),
+            fetch(nhwcIndex(z, y, x, 2))
+        );
+    }
+
+    /////////////////////////////////////////////////////////////////////////////////////////////////////////////
+    // Store, no broadcasting here. Assume we output full res gradient and then reduce using torch.sum outside
+    __device__ inline void store(unsigned int x, unsigned int y, unsigned int z, float _val)
+    {
+        store(_nhwcIndex(z, y, x, 0), _val);
+    }
+
+    __device__ inline void store(unsigned int x, unsigned int y, unsigned int z, vec3f _val)
+    {
+        store(_nhwcIndex(z, y, x, 0), _val.x);
+        store(_nhwcIndex(z, y, x, 1), _val.y);
+        store(_nhwcIndex(z, y, x, 2), _val.z);
+    }
+
+    /////////////////////////////////////////////////////////////////////////////////////////////////////////////
+    // Store gradient , no broadcasting here. Assume we output full res gradient and then reduce using torch.sum outside
+    __device__ inline void store_grad(unsigned int x, unsigned int y, unsigned int z, float _val)
+    {
+        store_grad(nhwcIndexContinuous(z, y, x, 0), _val);
+    }
+
+    __device__ inline void store_grad(unsigned int x, unsigned int y, unsigned int z, vec3f _val)
+    {
+        store_grad(nhwcIndexContinuous(z, y, x, 0), _val.x);
+        store_grad(nhwcIndexContinuous(z, y, x, 1), _val.y);
+        store_grad(nhwcIndexContinuous(z, y, x, 2), _val.z);
+    }
+#endif
+
+};

nvdiffmodeling/src/renderutils/c_src/torch_bindings.cpp

@@ -0,0 +1,793 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/CUDAUtils.h>
+#include <algorithm>
+#include <string>
+
+#define NVDR_CHECK_CUDA_ERROR(CUDA_CALL) { cudaError_t err = CUDA_CALL; AT_CUDA_CHECK(cudaGetLastError()); }
+#define NVDR_CHECK_GL_ERROR(GL_CALL) { GL_CALL; GLenum err = glGetError(); TORCH_CHECK(err == GL_NO_ERROR, "OpenGL error: ", getGLErrorString(err), "[", #GL_CALL, ";]"); }
+#define CHECK_TENSOR(X, DIMS, CHANNELS) \
+    TORCH_CHECK(X.is_cuda(), #X " must be a cuda tensor") \
+    TORCH_CHECK(X.scalar_type() == torch::kFloat || X.scalar_type() == torch::kBFloat16, #X " must be fp32 or bf16") \
+    TORCH_CHECK(X.dim() == DIMS, #X " must have " #DIMS " dimensions") \
+    TORCH_CHECK(X.size(DIMS - 1) == CHANNELS, #X " must have " #CHANNELS " channels")
+
+#include "common.h"
+#include "loss.h"
+#include "normal.h"
+#include "bsdf.h"
+#include "mesh.h"
+
+#define BLOCK_X 8
+#define BLOCK_Y 8
+
+//------------------------------------------------------------------------
+// mesh.cu
+
+void xfmPointsFwdKernel(XfmKernelParams p);
+void xfmPointsBwdKernel(XfmKernelParams p);
+
+//------------------------------------------------------------------------
+// loss.cu
+
+void imgLossFwdKernel(LossKernelParams p);
+void imgLossBwdKernel(LossKernelParams p);
+
+//------------------------------------------------------------------------
+// normal.cu
+
+void PrepareShadingNormalFwdKernel(PrepareShadingNormalKernelParams p);
+void PrepareShadingNormalBwdKernel(PrepareShadingNormalKernelParams p);
+
+//------------------------------------------------------------------------
+// bsdf.cu
+
+void LambertFwdKernel(LambertKernelParams p);
+void LambertBwdKernel(LambertKernelParams p);
+
+void FresnelShlickFwdKernel(FresnelShlickKernelParams p);
+void FresnelShlickBwdKernel(FresnelShlickKernelParams p);
+
+void ndfGGXFwdKernel(NdfGGXParams p);
+void ndfGGXBwdKernel(NdfGGXParams p);
+
+void lambdaGGXFwdKernel(NdfGGXParams p);
+void lambdaGGXBwdKernel(NdfGGXParams p);
+
+void maskingSmithFwdKernel(MaskingSmithParams p);
+void maskingSmithBwdKernel(MaskingSmithParams p);
+
+void pbrSpecularFwdKernel(PbrSpecular p);
+void pbrSpecularBwdKernel(PbrSpecular p);
+
+void pbrBSDFFwdKernel(PbrBSDF p);
+void pbrBSDFBwdKernel(PbrBSDF p);
+
+//------------------------------------------------------------------------
+// Tensor helpers
+
+void update_grid(dim3 &gridSize, torch::Tensor x)
+{
+    gridSize.x = std::max(gridSize.x, (uint32_t)x.size(2));
+    gridSize.y = std::max(gridSize.y, (uint32_t)x.size(1));
+    gridSize.z = std::max(gridSize.z, (uint32_t)x.size(0));
+}
+
+template<typename... Ts>
+void update_grid(dim3& gridSize, torch::Tensor x, Ts&&... vs)
+{
+    gridSize.x = std::max(gridSize.x, (uint32_t)x.size(2));
+    gridSize.y = std::max(gridSize.y, (uint32_t)x.size(1));
+    gridSize.z = std::max(gridSize.z, (uint32_t)x.size(0));
+    update_grid(gridSize, std::forward<Ts>(vs)...);
+}
+
+Tensor make_cuda_tensor(torch::Tensor val, dim3 outDims, torch::Tensor* grad = nullptr)
+{
+    Tensor res;
+    for (int i = 0; i < val.dim(); ++i)
+    {
+        res.dims[i] = val.size(i);
+        res.strides[i] = val.stride(i);
+    }
+
+    res.fp16 = val.scalar_type() == torch::kBFloat16;
+    res.val = res.fp16 ? (void*)val.data_ptr<torch::BFloat16>() : (void*)val.data_ptr<float>();
+    res.d_val = nullptr;
+    if (grad != nullptr)
+    {
+        if (val.dim() == 4)
+            *grad = torch::empty({ outDims.z, outDims.y, outDims.x, val.size(3) }, torch::TensorOptions().dtype(res.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA));
+        else // 3
+            *grad = torch::empty({ outDims.z, outDims.x, val.size(2) }, torch::TensorOptions().dtype(res.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA));
+
+        res.d_val = res.fp16 ? (void*)grad->data_ptr<torch::BFloat16>() : (void*)grad->data_ptr<float>();
+    }
+    return res;
+}
+
+//------------------------------------------------------------------------
+// prepare_shading_normal
+
+torch::Tensor prepare_shading_normal_fwd(torch::Tensor pos, torch::Tensor view_pos, torch::Tensor perturbed_nrm, torch::Tensor smooth_nrm, torch::Tensor smooth_tng, torch::Tensor geom_nrm, bool two_sided_shading, bool opengl, bool fp16)
+{
+    CHECK_TENSOR(pos, 4, 3);
+    CHECK_TENSOR(view_pos, 4, 3);
+    CHECK_TENSOR(perturbed_nrm, 4, 3);
+    CHECK_TENSOR(smooth_nrm, 4, 3);
+    CHECK_TENSOR(smooth_tng, 4, 3);
+    CHECK_TENSOR(geom_nrm, 4, 3);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    PrepareShadingNormalKernelParams p;
+    p.two_sided_shading = two_sided_shading;
+    p.opengl = opengl;
+    p.out.fp16 = fp16;
+    update_grid(p.gridSize, pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = torch::empty({ p.gridSize.z, p.gridSize.y, p.gridSize.x, 3 }, opts);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    // Setup tensors
+    p.pos = make_cuda_tensor(pos, p.gridSize);
+    p.view_pos = make_cuda_tensor(view_pos, p.gridSize);
+    p.perturbed_nrm = make_cuda_tensor(perturbed_nrm, p.gridSize);
+    p.smooth_nrm = make_cuda_tensor(smooth_nrm, p.gridSize);
+    p.smooth_tng = make_cuda_tensor(smooth_tng, p.gridSize);
+    p.geom_nrm = make_cuda_tensor(geom_nrm, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)PrepareShadingNormalFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> prepare_shading_normal_bwd(torch::Tensor pos, torch::Tensor view_pos, torch::Tensor perturbed_nrm, torch::Tensor smooth_nrm, torch::Tensor smooth_tng, torch::Tensor geom_nrm, torch::Tensor grad, bool two_sided_shading, bool opengl)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    PrepareShadingNormalKernelParams p;
+    p.two_sided_shading = two_sided_shading;
+    p.opengl = opengl;
+    update_grid(p.gridSize, pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    // Setup tensors
+    torch::Tensor pos_grad, view_pos_grad, perturbed_nrm_grad, smooth_nrm_grad, smooth_tng_grad, geom_nrm_grad;
+    p.pos = make_cuda_tensor(pos, p.gridSize, &pos_grad);
+    p.view_pos = make_cuda_tensor(view_pos, p.gridSize, &view_pos_grad);
+    p.perturbed_nrm = make_cuda_tensor(perturbed_nrm, p.gridSize, &perturbed_nrm_grad);
+    p.smooth_nrm = make_cuda_tensor(smooth_nrm, p.gridSize, &smooth_nrm_grad);
+    p.smooth_tng = make_cuda_tensor(smooth_tng, p.gridSize, &smooth_tng_grad);
+    p.geom_nrm = make_cuda_tensor(geom_nrm, p.gridSize, &geom_nrm_grad);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)PrepareShadingNormalBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>(pos_grad, view_pos_grad, perturbed_nrm_grad, smooth_nrm_grad, smooth_tng_grad, geom_nrm_grad);
+}
+
+//------------------------------------------------------------------------
+// lambert
+
+torch::Tensor lambert_fwd(torch::Tensor nrm, torch::Tensor wi, bool fp16)
+{
+    CHECK_TENSOR(nrm, 4, 3);
+    CHECK_TENSOR(wi, 4, 3);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    LambertKernelParams p;
+    p.out.fp16 = fp16;
+    update_grid(p.gridSize, nrm, wi);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = torch::empty({ p.gridSize.z, p.gridSize.y, p.gridSize.x, 1 }, opts);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    p.nrm = make_cuda_tensor(nrm, p.gridSize);
+    p.wi = make_cuda_tensor(wi, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)LambertFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> lambert_bwd(torch::Tensor nrm, torch::Tensor wi, torch::Tensor grad)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    LambertKernelParams p;
+    update_grid(p.gridSize, nrm, wi);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    torch::Tensor nrm_grad, wi_grad;
+    p.nrm = make_cuda_tensor(nrm, p.gridSize, &nrm_grad);
+    p.wi = make_cuda_tensor(wi, p.gridSize, &wi_grad);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)LambertBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return std::tuple<torch::Tensor, torch::Tensor>(nrm_grad, wi_grad);
+}
+
+//------------------------------------------------------------------------
+// fresnel_shlick
+
+torch::Tensor fresnel_shlick_fwd(torch::Tensor f0, torch::Tensor f90, torch::Tensor cosTheta, bool fp16)
+{
+    CHECK_TENSOR(f0, 4, 3);
+    CHECK_TENSOR(f90, 4, 3);
+    CHECK_TENSOR(cosTheta, 4, 1);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    FresnelShlickKernelParams p;
+    p.out.fp16 = fp16;
+    update_grid(p.gridSize, f0, f90, cosTheta);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = torch::empty({ p.gridSize.z, p.gridSize.y, p.gridSize.x, 3 }, opts);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    p.f0 = make_cuda_tensor(f0, p.gridSize);
+    p.f90 = make_cuda_tensor(f90, p.gridSize);
+    p.cosTheta = make_cuda_tensor(cosTheta, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)FresnelShlickFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> fresnel_shlick_bwd(torch::Tensor f0, torch::Tensor f90, torch::Tensor cosTheta, torch::Tensor grad)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    FresnelShlickKernelParams p;
+    update_grid(p.gridSize, f0, f90, cosTheta);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    torch::Tensor f0_grad, f90_grad, cosT_grad;
+    p.f0 = make_cuda_tensor(f0, p.gridSize, &f0_grad);
+    p.f90 = make_cuda_tensor(f90, p.gridSize, &f90_grad);
+    p.cosTheta = make_cuda_tensor(cosTheta, p.gridSize, &cosT_grad);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)FresnelShlickBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>(f0_grad, f90_grad, cosT_grad);
+}
+
+//------------------------------------------------------------------------
+// ndf_ggd
+
+torch::Tensor ndf_ggx_fwd(torch::Tensor alphaSqr, torch::Tensor cosTheta, bool fp16)
+{
+    CHECK_TENSOR(alphaSqr, 4, 1);
+    CHECK_TENSOR(cosTheta, 4, 1);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    NdfGGXParams p;
+    p.out.fp16 = fp16;
+    update_grid(p.gridSize, alphaSqr, cosTheta);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = torch::empty({ p.gridSize.z, p.gridSize.y, p.gridSize.x, 1 }, opts);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    p.alphaSqr = make_cuda_tensor(alphaSqr, p.gridSize);
+    p.cosTheta = make_cuda_tensor(cosTheta, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)ndfGGXFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> ndf_ggx_bwd(torch::Tensor alphaSqr, torch::Tensor cosTheta, torch::Tensor grad)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    NdfGGXParams p;
+    update_grid(p.gridSize, alphaSqr, cosTheta);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    torch::Tensor alphaSqr_grad, cosTheta_grad;
+    p.alphaSqr = make_cuda_tensor(alphaSqr, p.gridSize, &alphaSqr_grad);
+    p.cosTheta = make_cuda_tensor(cosTheta, p.gridSize, &cosTheta_grad);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)ndfGGXBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return std::tuple<torch::Tensor, torch::Tensor>(alphaSqr_grad, cosTheta_grad);
+}
+
+//------------------------------------------------------------------------
+// lambda_ggx
+
+torch::Tensor lambda_ggx_fwd(torch::Tensor alphaSqr, torch::Tensor cosTheta, bool fp16)
+{
+    CHECK_TENSOR(alphaSqr, 4, 1);
+    CHECK_TENSOR(cosTheta, 4, 1);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    NdfGGXParams p;
+    p.out.fp16 = fp16;
+    update_grid(p.gridSize, alphaSqr, cosTheta);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = torch::empty({ p.gridSize.z, p.gridSize.y, p.gridSize.x, 1 }, opts);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    p.alphaSqr = make_cuda_tensor(alphaSqr, p.gridSize);
+    p.cosTheta = make_cuda_tensor(cosTheta, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)lambdaGGXFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> lambda_ggx_bwd(torch::Tensor alphaSqr, torch::Tensor cosTheta, torch::Tensor grad)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    NdfGGXParams p;
+    update_grid(p.gridSize, alphaSqr, cosTheta);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    torch::Tensor alphaSqr_grad, cosTheta_grad;
+    p.alphaSqr = make_cuda_tensor(alphaSqr, p.gridSize, &alphaSqr_grad);
+    p.cosTheta = make_cuda_tensor(cosTheta, p.gridSize, &cosTheta_grad);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)lambdaGGXBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return std::tuple<torch::Tensor, torch::Tensor>(alphaSqr_grad, cosTheta_grad);
+}
+
+//------------------------------------------------------------------------
+// masking_smith
+
+torch::Tensor masking_smith_fwd(torch::Tensor alphaSqr, torch::Tensor cosThetaI, torch::Tensor cosThetaO, bool fp16)
+{
+    CHECK_TENSOR(alphaSqr, 4, 1);
+    CHECK_TENSOR(cosThetaI, 4, 1);
+    CHECK_TENSOR(cosThetaO, 4, 1);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    MaskingSmithParams p;
+    p.out.fp16 = fp16;
+    update_grid(p.gridSize, alphaSqr, cosThetaI, cosThetaO);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = torch::empty({ p.gridSize.z, p.gridSize.y, p.gridSize.x, 1 }, opts);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    p.alphaSqr = make_cuda_tensor(alphaSqr, p.gridSize);
+    p.cosThetaI = make_cuda_tensor(cosThetaI, p.gridSize);
+    p.cosThetaO = make_cuda_tensor(cosThetaO, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)maskingSmithFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> masking_smith_bwd(torch::Tensor alphaSqr, torch::Tensor cosThetaI, torch::Tensor cosThetaO, torch::Tensor grad)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    MaskingSmithParams p;
+    update_grid(p.gridSize, alphaSqr, cosThetaI, cosThetaO);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    torch::Tensor alphaSqr_grad, cosThetaI_grad, cosThetaO_grad;
+    p.alphaSqr = make_cuda_tensor(alphaSqr, p.gridSize, &alphaSqr_grad);
+    p.cosThetaI = make_cuda_tensor(cosThetaI, p.gridSize, &cosThetaI_grad);
+    p.cosThetaO = make_cuda_tensor(cosThetaO, p.gridSize, &cosThetaO_grad);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)maskingSmithBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>(alphaSqr_grad, cosThetaI_grad, cosThetaO_grad);
+}
+
+//------------------------------------------------------------------------
+// pbr_specular
+
+torch::Tensor pbr_specular_fwd(torch::Tensor col, torch::Tensor nrm, torch::Tensor wo, torch::Tensor wi, torch::Tensor alpha, float min_roughness, bool fp16)
+{
+    CHECK_TENSOR(col, 4, 3);
+    CHECK_TENSOR(nrm, 4, 3);
+    CHECK_TENSOR(wo, 4, 3);
+    CHECK_TENSOR(wi, 4, 3);
+    CHECK_TENSOR(alpha, 4, 1);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    PbrSpecular p;
+    p.out.fp16 = fp16;
+    p.min_roughness = min_roughness;
+    update_grid(p.gridSize, col, nrm, wo, wi, alpha);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = torch::empty({ p.gridSize.z, p.gridSize.y, p.gridSize.x, 3 }, opts);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    p.col = make_cuda_tensor(col, p.gridSize);
+    p.nrm = make_cuda_tensor(nrm, p.gridSize);
+    p.wo = make_cuda_tensor(wo, p.gridSize);
+    p.wi = make_cuda_tensor(wi, p.gridSize);
+    p.alpha = make_cuda_tensor(alpha, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)pbrSpecularFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> pbr_specular_bwd(torch::Tensor col, torch::Tensor nrm, torch::Tensor wo, torch::Tensor wi, torch::Tensor alpha, float min_roughness, torch::Tensor grad)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    PbrSpecular p;
+    update_grid(p.gridSize, col, nrm, wo, wi, alpha);
+    p.min_roughness = min_roughness;
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    torch::Tensor col_grad, nrm_grad, wo_grad, wi_grad, alpha_grad;
+    p.col = make_cuda_tensor(col, p.gridSize, &col_grad);
+    p.nrm = make_cuda_tensor(nrm, p.gridSize, &nrm_grad);
+    p.wo = make_cuda_tensor(wo, p.gridSize, &wo_grad);
+    p.wi = make_cuda_tensor(wi, p.gridSize, &wi_grad);
+    p.alpha = make_cuda_tensor(alpha, p.gridSize, &alpha_grad);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)pbrSpecularBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>(col_grad, nrm_grad, wo_grad, wi_grad, alpha_grad);
+}
+
+//------------------------------------------------------------------------
+// pbr_bsdf
+
+torch::Tensor pbr_bsdf_fwd(torch::Tensor kd, torch::Tensor arm, torch::Tensor pos, torch::Tensor nrm, torch::Tensor view_pos, torch::Tensor light_pos, float min_roughness, bool fp16)
+{
+    CHECK_TENSOR(kd, 4, 3);
+    CHECK_TENSOR(arm, 4, 3);
+    CHECK_TENSOR(pos, 4, 3);
+    CHECK_TENSOR(nrm, 4, 3);
+    CHECK_TENSOR(view_pos, 4, 3);
+    CHECK_TENSOR(light_pos, 4, 3);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    PbrBSDF p;
+    p.out.fp16 = fp16;
+    p.min_roughness = min_roughness;
+    update_grid(p.gridSize, kd, arm, pos, nrm, view_pos, light_pos);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = torch::empty({ p.gridSize.z, p.gridSize.y, p.gridSize.x, 3 }, opts);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    p.kd = make_cuda_tensor(kd, p.gridSize);
+    p.arm = make_cuda_tensor(arm, p.gridSize);
+    p.pos = make_cuda_tensor(pos, p.gridSize);
+    p.nrm = make_cuda_tensor(nrm, p.gridSize);
+    p.view_pos = make_cuda_tensor(view_pos, p.gridSize);
+    p.light_pos = make_cuda_tensor(light_pos, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)pbrBSDFFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> pbr_bsdf_bwd(torch::Tensor kd, torch::Tensor arm, torch::Tensor pos, torch::Tensor nrm, torch::Tensor view_pos, torch::Tensor light_pos, float min_roughness, torch::Tensor grad)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    PbrBSDF p;
+    update_grid(p.gridSize, kd, arm, pos, nrm, view_pos, light_pos);
+    p.min_roughness = min_roughness;
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    torch::Tensor kd_grad, arm_grad, pos_grad, nrm_grad, view_pos_grad, light_pos_grad;
+    p.kd = make_cuda_tensor(kd, p.gridSize, &kd_grad);
+    p.arm = make_cuda_tensor(arm, p.gridSize, &arm_grad);
+    p.pos = make_cuda_tensor(pos, p.gridSize, &pos_grad);
+    p.nrm = make_cuda_tensor(nrm, p.gridSize, &nrm_grad);
+    p.view_pos = make_cuda_tensor(view_pos, p.gridSize, &view_pos_grad);
+    p.light_pos = make_cuda_tensor(light_pos, p.gridSize, &light_pos_grad);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)pbrBSDFBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>(kd_grad, arm_grad, pos_grad, nrm_grad, view_pos_grad, light_pos_grad);
+}
+
+//------------------------------------------------------------------------
+// loss function
+
+LossType strToLoss(std::string str)
+{
+    if (str == "mse")
+        return LOSS_MSE;
+    else if (str == "relmse")
+        return LOSS_RELMSE;
+    else if (str == "smape")
+        return LOSS_SMAPE;
+    else
+        return LOSS_L1;
+}
+
+torch::Tensor image_loss_fwd(torch::Tensor img, torch::Tensor target, std::string loss, std::string tonemapper, bool fp16)
+{
+    CHECK_TENSOR(img, 4, 3);
+    CHECK_TENSOR(target, 4, 3);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    LossKernelParams p;
+    p.out.fp16 = fp16;
+    p.loss = strToLoss(loss);
+    p.tonemapper = tonemapper == "log_srgb" ? TONEMAPPER_LOG_SRGB : TONEMAPPER_NONE;
+    update_grid(p.gridSize, img, target);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 warpSize = getWarpSize(blockSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = torch::empty({ (p.gridSize.z - 1) / warpSize.z + 1, (p.gridSize.y - 1) / warpSize.y + 1, (p.gridSize.x - 1) / warpSize.x + 1, 1 }, opts);
+
+    p.img = make_cuda_tensor(img, p.gridSize);
+    p.target = make_cuda_tensor(target, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)imgLossFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> image_loss_bwd(torch::Tensor img, torch::Tensor target, torch::Tensor grad, std::string loss, std::string tonemapper)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    LossKernelParams p;
+    p.loss = strToLoss(loss);
+    p.tonemapper = tonemapper == "log_srgb" ? TONEMAPPER_LOG_SRGB : TONEMAPPER_NONE;
+    update_grid(p.gridSize, img, target);
+
+    // Choose launch parameters.
+    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
+    dim3 warpSize = getWarpSize(blockSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    torch::Tensor img_grad, target_grad;
+    p.img = make_cuda_tensor(img, p.gridSize, &img_grad);
+    p.target = make_cuda_tensor(target, p.gridSize, &target_grad);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)imgLossBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return std::tuple<torch::Tensor, torch::Tensor>(img_grad, target_grad);
+}
+
+//------------------------------------------------------------------------
+// transform function
+
+torch::Tensor xfm_fwd(torch::Tensor points, torch::Tensor matrix, bool isPoints, bool fp16)
+{
+    CHECK_TENSOR(points, 3, 3);
+    CHECK_TENSOR(matrix, 3, 4);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    XfmKernelParams p;
+    p.out.fp16 = fp16;
+    p.isPoints = isPoints;
+    p.gridSize.x = points.size(1);
+    p.gridSize.y = 1;
+    p.gridSize.z = std::max(matrix.size(0), points.size(0));
+
+    // Choose launch parameters.
+    dim3 blockSize(BLOCK_X * BLOCK_Y, 1, 1);
+    dim3 warpSize = getWarpSize(blockSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    // Allocate output tensors.
+    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
+    torch::Tensor out = isPoints ? torch::empty({ matrix.size(0), points.size(1), 4 }, opts) : torch::empty({ matrix.size(0), points.size(1), 3 }, opts);
+
+    p.points = make_cuda_tensor(points, p.gridSize);
+    p.matrix = make_cuda_tensor(matrix, p.gridSize);
+    p.out = make_cuda_tensor(out, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)xfmPointsFwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return out;
+}
+
+torch::Tensor xfm_bwd(torch::Tensor points, torch::Tensor matrix, torch::Tensor grad, bool isPoints)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    // Extract input parameters.
+    XfmKernelParams p;
+    p.isPoints = isPoints;
+    p.gridSize.x = points.size(1);
+    p.gridSize.y = 1;
+    p.gridSize.z = std::max(matrix.size(0), points.size(0));
+
+    // Choose launch parameters.
+    dim3 blockSize(BLOCK_X * BLOCK_Y, 1, 1);
+    dim3 warpSize = getWarpSize(blockSize);
+    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);
+
+    torch::Tensor points_grad;
+    p.points = make_cuda_tensor(points, p.gridSize, &points_grad);
+    p.matrix = make_cuda_tensor(matrix, p.gridSize);
+    p.out = make_cuda_tensor(grad, p.gridSize);
+
+    // Launch CUDA kernel.
+    void* args[] = { &p };
+    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)xfmPointsBwdKernel, gridSize, blockSize, args, 0, stream));
+
+    return points_grad;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("prepare_shading_normal_fwd", &prepare_shading_normal_fwd, "prepare_shading_normal_fwd");
+    m.def("prepare_shading_normal_bwd", &prepare_shading_normal_bwd, "prepare_shading_normal_bwd");
+    m.def("lambert_fwd", &lambert_fwd, "lambert_fwd");
+    m.def("lambert_bwd", &lambert_bwd, "lambert_bwd");
+    m.def("fresnel_shlick_fwd", &fresnel_shlick_fwd, "fresnel_shlick_fwd");
+    m.def("fresnel_shlick_bwd", &fresnel_shlick_bwd, "fresnel_shlick_bwd");
+    m.def("ndf_ggx_fwd", &ndf_ggx_fwd, "ndf_ggx_fwd");
+    m.def("ndf_ggx_bwd", &ndf_ggx_bwd, "ndf_ggx_bwd");
+    m.def("lambda_ggx_fwd", &lambda_ggx_fwd, "lambda_ggx_fwd");
+    m.def("lambda_ggx_bwd", &lambda_ggx_bwd, "lambda_ggx_bwd");
+    m.def("masking_smith_fwd", &masking_smith_fwd, "masking_smith_fwd");
+    m.def("masking_smith_bwd", &masking_smith_bwd, "masking_smith_bwd");
+    m.def("pbr_specular_fwd", &pbr_specular_fwd, "pbr_specular_fwd");
+    m.def("pbr_specular_bwd", &pbr_specular_bwd, "pbr_specular_bwd");
+    m.def("pbr_bsdf_fwd", &pbr_bsdf_fwd, "pbr_bsdf_fwd");
+    m.def("pbr_bsdf_bwd", &pbr_bsdf_bwd, "pbr_bsdf_bwd");
+    m.def("image_loss_fwd", &image_loss_fwd, "image_loss_fwd");
+    m.def("image_loss_bwd", &image_loss_bwd, "image_loss_bwd");
+    m.def("xfm_fwd", &xfm_fwd, "xfm_fwd");
+    m.def("xfm_bwd", &xfm_bwd, "xfm_bwd");
+}

nvdiffmodeling/src/renderutils/c_src/vec3f.h

@@ -0,0 +1,106 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once 
+
+struct vec3f
+{
+    float x, y, z;
+
+#ifdef __CUDACC__
+    __device__ vec3f() { }
+    __device__ vec3f(float v) { x = v; y = v; z = v; }
+    __device__ vec3f(float _x, float _y, float _z) { x = _x; y = _y; z = _z; }
+    __device__ vec3f(float3 v) { x = v.x; y = v.y; z = v.z; }
+
+    __device__ inline vec3f& operator+=(const vec3f& b) { x += b.x; y += b.y; z += b.z; return *this; }
+    __device__ inline vec3f& operator-=(const vec3f& b) { x -= b.x; y -= b.y; z -= b.z; return *this; }
+    __device__ inline vec3f& operator*=(const vec3f& b) { x *= b.x; y *= b.y; z *= b.z; return *this; }
+    __device__ inline vec3f& operator/=(const vec3f& b) { x /= b.x; y /= b.y; z /= b.z; return *this; }
+#endif
+};
+
+#ifdef __CUDACC__
+__device__ static inline vec3f operator+(const vec3f& a, const vec3f& b) { return vec3f(a.x + b.x, a.y + b.y, a.z + b.z); }
+__device__ static inline vec3f operator-(const vec3f& a, const vec3f& b) { return vec3f(a.x - b.x, a.y - b.y, a.z - b.z); }
+__device__ static inline vec3f operator*(const vec3f& a, const vec3f& b) { return vec3f(a.x * b.x, a.y * b.y, a.z * b.z); }
+__device__ static inline vec3f operator/(const vec3f& a, const vec3f& b) { return vec3f(a.x / b.x, a.y / b.y, a.z / b.z); }
+__device__ static inline vec3f operator-(const vec3f& a) { return vec3f(-a.x, -a.y, -a.z); }
+
+__device__ static inline float sum(vec3f a)
+{
+    return a.x + a.y + a.z;
+}
+
+__device__ static inline vec3f cross(vec3f a, vec3f b)
+{
+    vec3f out;
+    out.x = a.y * b.z - a.z * b.y;
+    out.y = a.z * b.x - a.x * b.z;
+    out.z = a.x * b.y - a.y * b.x;
+    return out;
+}
+
+__device__ static inline void bwdCross(vec3f a, vec3f b, vec3f &d_a, vec3f &d_b, vec3f d_out)
+{
+    d_a.x += d_out.z * b.y - d_out.y * b.z;
+    d_a.y += d_out.x * b.z - d_out.z * b.x;
+    d_a.z += d_out.y * b.x - d_out.x * b.y;
+
+    d_b.x += d_out.y * a.z - d_out.z * a.y;
+    d_b.y += d_out.z * a.x - d_out.x * a.z;
+    d_b.z += d_out.x * a.y - d_out.y * a.x;
+}
+
+__device__ static inline float dot(vec3f a, vec3f b)
+{
+    return a.x * b.x + a.y * b.y + a.z * b.z;
+}
+
+__device__ static inline void bwdDot(vec3f a, vec3f b, vec3f& d_a, vec3f& d_b, float d_out)
+{
+    d_a.x += d_out * b.x; d_a.y += d_out * b.y; d_a.z += d_out * b.z;
+    d_b.x += d_out * a.x; d_b.y += d_out * a.y; d_b.z += d_out * a.z;
+}
+
+__device__ static inline vec3f reflect(vec3f x, vec3f n)
+{
+    return n * 2.0f * dot(n, x) - x;
+}
+
+__device__ static inline void bwdReflect(vec3f x, vec3f n, vec3f& d_x, vec3f& d_n, const vec3f d_out)
+{
+    d_x.x += d_out.x * (2 * n.x * n.x - 1) + d_out.y * (2 * n.x * n.y) + d_out.z * (2 * n.x * n.z);
+    d_x.y += d_out.x * (2 * n.x * n.y) + d_out.y * (2 * n.y * n.y - 1) + d_out.z * (2 * n.y * n.z);
+    d_x.z += d_out.x * (2 * n.x * n.z) + d_out.y * (2 * n.y * n.z) + d_out.z * (2 * n.z * n.z - 1);
+
+    d_n.x += d_out.x * (2 * (2 * n.x * x.x + n.y * x.y + n.z * x.z)) + d_out.y * (2 * n.y * x.x) + d_out.z * (2 * n.z * x.x);
+    d_n.y += d_out.x * (2 * n.x * x.y) + d_out.y * (2 * (n.x * x.x + 2 * n.y * x.y + n.z * x.z)) + d_out.z * (2 * n.z * x.y);
+    d_n.z += d_out.x * (2 * n.x * x.z) + d_out.y * (2 * n.y * x.z) + d_out.z * (2 * (n.x * x.x + n.y * x.y + 2 * n.z * x.z));
+}
+
+__device__ static inline vec3f safeNormalize(vec3f v)
+{
+    float l = sqrtf(v.x * v.x + v.y * v.y + v.z * v.z);
+    return l > 0.0f ? (v / l) : vec3f(0.0f);
+}
+
+__device__ static inline void bwdSafeNormalize(const vec3f v, vec3f& d_v, const vec3f d_out)
+{
+
+    float l = sqrtf(v.x * v.x + v.y * v.y + v.z * v.z);
+    if (l > 0.0f)
+    {
+        float fac = 1.0 / powf(v.x * v.x + v.y * v.y + v.z * v.z, 1.5f);
+        d_v.x += (d_out.x * (v.y * v.y + v.z * v.z) - d_out.y * (v.x * v.y) - d_out.z * (v.x * v.z)) * fac;
+        d_v.y += (d_out.y * (v.x * v.x + v.z * v.z) - d_out.x * (v.y * v.x) - d_out.z * (v.y * v.z)) * fac;
+        d_v.z += (d_out.z * (v.x * v.x + v.y * v.y) - d_out.x * (v.z * v.x) - d_out.y * (v.z * v.y)) * fac;
+    }
+}
+
+#endif

nvdiffmodeling/src/renderutils/c_src/vec4f.h

@@ -0,0 +1,22 @@
+// Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto. Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once 
+
+struct vec4f
+{
+    float x, y, z, w;
+
+#ifdef __CUDACC__
+    __device__ vec4f() { }
+    __device__ vec4f(float v) { x = v; y = v; z = v; w = v; }
+    __device__ vec4f(float _x, float _y, float _z, float _w) { x = _x; y = _y; z = _z; w = _w; }
+    __device__ vec4f(float4 v) { x = v.x; y = v.y; z = v.z; w = v.w; }
+#endif
+};
+

nvdiffmodeling/src/renderutils/loss.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import torch
+
+#----------------------------------------------------------------------------
+# HDR image losses
+#----------------------------------------------------------------------------
+
+def _tonemap_srgb(f):
+    return torch.where(f > 0.0031308, torch.pow(torch.clamp(f, min=0.0031308), 1.0/2.4)*1.055 - 0.055, 12.92*f)
+
+def _SMAPE(img, target, eps=0.01):
+    nom = torch.abs(img - target)
+    denom = torch.abs(img) + torch.abs(target) + 0.01
+    return torch.mean(nom / denom)
+
+def _RELMSE(img, target, eps=0.1):
+    nom = (img - target) * (img - target)
+    denom = img * img + target * target + 0.1 
+    return torch.mean(nom / denom)
+
+def image_loss_fn(img, target, loss, tonemapper):
+    if tonemapper == 'log_srgb':
+        img    = _tonemap_srgb(torch.log(torch.clamp(img, min=0, max=65535) + 1))
+        target = _tonemap_srgb(torch.log(torch.clamp(target, min=0, max=65535) + 1))
+
+    if loss == 'mse':
+        return torch.nn.functional.mse_loss(img, target)
+    elif loss == 'smape':
+        return _SMAPE(img, target)
+    elif loss == 'relmse':
+        return _RELMSE(img, target)
+    else:
+        return torch.nn.functional.l1_loss(img, target)

nvdiffmodeling/src/renderutils/ops.py

@@ -0,0 +1,425 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import numpy as np
+import os
+import sys
+import torch
+import torch.utils.cpp_extension
+
+from .bsdf import *
+from .loss import *
+
+#----------------------------------------------------------------------------
+# C++/Cuda plugin compiler/loader.
+
+_plugin = None
+if _plugin is None:
+
+    # Make sure we can find the necessary compiler and libary binaries.
+    if os.name == 'nt':
+        def find_cl_path():
+            import glob
+            for edition in ['Enterprise', 'Professional', 'BuildTools', 'Community']:
+                paths = sorted(glob.glob(r"C:\Program Files (x86)\Microsoft Visual Studio\*\%s\VC\Tools\MSVC\*\bin\Hostx64\x64" % edition), reverse=True)
+                if paths:
+                    return paths[0]
+
+        # If cl.exe is not on path, try to find it.
+        if os.system("where cl.exe >nul 2>nul") != 0:
+            cl_path = find_cl_path()
+            if cl_path is None:
+                raise RuntimeError("Could not locate a supported Microsoft Visual C++ installation")
+            os.environ['PATH'] += ';' + cl_path
+
+    # Linker options.
+    if os.name == 'posix':
+        ldflags = ['-lcuda']
+    elif os.name == 'nt':
+        ldflags = ['/DEFAULTLIB:cuda']
+
+    # List of sources.
+    source_files = [
+        'c_src/mesh.cu',
+        'c_src/loss.cu',
+        'c_src/bsdf.cu',
+        'c_src/normal.cu',
+        'c_src/common.cpp',
+        'c_src/torch_bindings.cpp'
+    ]
+
+    # Some containers set this to contain old architectures that won't compile. We only need the one installed in the machine.
+    os.environ['TORCH_CUDA_ARCH_LIST'] = ''
+
+    # Compile and load.
+    source_paths = [os.path.join(os.path.dirname(__file__), fn) for fn in source_files]
+    torch.utils.cpp_extension.load(name='renderutils_plugin', sources=source_paths, extra_ldflags=ldflags, with_cuda=True, verbose=True)
+
+    # Import, cache, and return the compiled module.
+    import renderutils_plugin
+    _plugin = renderutils_plugin
+
+#----------------------------------------------------------------------------
+# Internal kernels, just used for testing functionality
+
+class _fresnel_shlick_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, f0, f90, cosTheta):
+        out = _plugin.fresnel_shlick_fwd(f0, f90, cosTheta, False)
+        ctx.save_for_backward(f0, f90, cosTheta)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        f0, f90, cosTheta = ctx.saved_variables
+        return _plugin.fresnel_shlick_bwd(f0, f90, cosTheta, dout) + (None,)
+
+def _fresnel_shlick(f0, f90, cosTheta, use_python=False):
+    if use_python:
+        out = bsdf_fresnel_shlick(f0, f90, cosTheta)
+    else:
+        out = _fresnel_shlick_func.apply(f0, f90, cosTheta)
+
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of _fresnel_shlick contains inf or NaN"
+    return out
+
+
+class _ndf_ggx_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, alphaSqr, cosTheta):
+        out = _plugin.ndf_ggx_fwd(alphaSqr, cosTheta, False)
+        ctx.save_for_backward(alphaSqr, cosTheta)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        alphaSqr, cosTheta = ctx.saved_variables
+        return _plugin.ndf_ggx_bwd(alphaSqr, cosTheta, dout) + (None,)
+
+def _ndf_ggx(alphaSqr, cosTheta, use_python=False):
+    if use_python:
+        out = bsdf_ndf_ggx(alphaSqr, cosTheta)
+    else:
+        out = _ndf_ggx_func.apply(alphaSqr, cosTheta)
+
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of _ndf_ggx contains inf or NaN"
+    return out
+
+class _lambda_ggx_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, alphaSqr, cosTheta):
+        out = _plugin.lambda_ggx_fwd(alphaSqr, cosTheta, False)
+        ctx.save_for_backward(alphaSqr, cosTheta)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        alphaSqr, cosTheta = ctx.saved_variables
+        return _plugin.lambda_ggx_bwd(alphaSqr, cosTheta, dout) + (None,)
+
+def _lambda_ggx(alphaSqr, cosTheta, use_python=False):
+    if use_python:
+        out = bsdf_lambda_ggx(alphaSqr, cosTheta)
+    else:
+        out = _lambda_ggx_func.apply(alphaSqr, cosTheta)
+
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of _lambda_ggx contains inf or NaN"
+    return out
+
+class _masking_smith_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, alphaSqr, cosThetaI, cosThetaO):
+        ctx.save_for_backward(alphaSqr, cosThetaI, cosThetaO)
+        out = _plugin.masking_smith_fwd(alphaSqr, cosThetaI, cosThetaO, False)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        alphaSqr, cosThetaI, cosThetaO = ctx.saved_variables
+        return _plugin.masking_smith_bwd(alphaSqr, cosThetaI, cosThetaO, dout) + (None,)
+
+def _masking_smith(alphaSqr, cosThetaI, cosThetaO, use_python=False):
+    if use_python:
+        out = bsdf_masking_smith_ggx_correlated(alphaSqr, cosThetaI, cosThetaO)
+    else:
+        out = _masking_smith_func.apply(alphaSqr, cosThetaI, cosThetaO)
+
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of _masking_smith contains inf or NaN"
+    return out
+
+#----------------------------------------------------------------------------
+# Shading normal setup (bump mapping + bent normals)
+
+class _prepare_shading_normal_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm, two_sided_shading, opengl):
+        ctx.two_sided_shading, ctx.opengl = two_sided_shading, opengl
+        out = _plugin.prepare_shading_normal_fwd(pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm, two_sided_shading, opengl, False)
+        ctx.save_for_backward(pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm = ctx.saved_variables
+        return _plugin.prepare_shading_normal_bwd(pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm, dout, ctx.two_sided_shading, ctx.opengl) + (None, None, None)
+
+def prepare_shading_normal(pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm, two_sided_shading=True, opengl=True, use_python=False):
+    '''Takes care of all corner cases and produces a final normal used for shading:
+        - Constructs tangent space
+        - Flips normal direction based on geometric normal for two sided Shading
+        - Perturbs shading normal by normal map
+        - Bends backfacing normals towards the camera to avoid shading artifacts
+
+        All tensors assume a shape of [minibatch_size, height, width, 3] or broadcastable equivalent.
+
+    Args:
+        pos: World space g-buffer position.
+        view_pos: Camera position in world space (typically using broadcasting).
+        perturbed_nrm: Trangent-space normal perturbation from normal map lookup.
+        smooth_nrm: Interpolated vertex normals.
+        smooth_tng: Interpolated vertex tangents.
+        geom_nrm: Geometric (face) normals.
+        two_sided_shading: Use one/two sided shading
+        opengl: Use OpenGL/DirectX normal map conventions 
+        use_python: Use PyTorch implementation (for validation)
+    Returns:
+        Final shading normal
+    '''    
+
+    if perturbed_nrm is None:
+        perturbed_nrm = torch.tensor([0, 0, 1], dtype=torch.float32, device='cuda', requires_grad=False)[None, None, None, ...]
+    
+    if use_python:
+        out = bsdf_prepare_shading_normal(pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm, two_sided_shading, opengl)
+    else:
+        out = _prepare_shading_normal_func.apply(pos, view_pos, perturbed_nrm, smooth_nrm, smooth_tng, geom_nrm, two_sided_shading, opengl)
+    
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of prepare_shading_normal contains inf or NaN"
+    return out
+
+#----------------------------------------------------------------------------
+# BSDF functions
+
+class _lambert_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, nrm, wi):
+        out = _plugin.lambert_fwd(nrm, wi, False)
+        ctx.save_for_backward(nrm, wi)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        nrm, wi = ctx.saved_variables
+        return _plugin.lambert_bwd(nrm, wi, dout) + (None,)
+
+def lambert(nrm, wi, use_python=False):
+    '''Lambertian bsdf. 
+    All tensors assume a shape of [minibatch_size, height, width, 3] or broadcastable equivalent.
+
+    Args:
+        nrm: World space shading normal.
+        wi: World space light vector.
+        use_python: Use PyTorch implementation (for validation)
+
+    Returns:
+        Shaded diffuse value with shape [minibatch_size, height, width, 1]
+    '''
+
+    if use_python:
+        out = bsdf_lambert(nrm, wi)
+    else:
+        out = _lambert_func.apply(nrm, wi)
+ 
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of lambert contains inf or NaN"
+    return out
+
+class _pbr_specular_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, col, nrm, wo, wi, alpha, min_roughness):
+        ctx.save_for_backward(col, nrm, wo, wi, alpha)
+        ctx.min_roughness = min_roughness
+        out = _plugin.pbr_specular_fwd(col, nrm, wo, wi, alpha, min_roughness, False)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        col, nrm, wo, wi, alpha = ctx.saved_variables
+        return _plugin.pbr_specular_bwd(col, nrm, wo, wi, alpha, ctx.min_roughness, dout) + (None, None)
+
+def pbr_specular(col, nrm, wo, wi, alpha, min_roughness=0.08, use_python=False):
+    '''Physically-based specular bsdf.
+    All tensors assume a shape of [minibatch_size, height, width, 3] or broadcastable equivalent unless otherwise noted.
+
+    Args:
+        col: Specular lobe color
+        nrm: World space shading normal.
+        wo: World space camera vector.
+        wi: World space light vector
+        alpha: Specular roughness parameter with shape [minibatch_size, height, width, 1]
+        min_roughness: Scalar roughness clamping threshold
+
+        use_python: Use PyTorch implementation (for validation)
+    Returns:
+        Shaded specular color
+    '''
+
+    if use_python:
+        out = bsdf_pbr_specular(col, nrm, wo, wi, alpha, min_roughness=min_roughness)
+    else:
+        out = _pbr_specular_func.apply(col, nrm, wo, wi, alpha, min_roughness)
+    
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of pbr_specular contains inf or NaN"
+    return out
+
+class _pbr_bsdf_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, kd, arm, pos, nrm, view_pos, light_pos, min_roughness):
+        ctx.save_for_backward(kd, arm, pos, nrm, view_pos, light_pos)
+        ctx.min_roughness = min_roughness
+        out = _plugin.pbr_bsdf_fwd(kd, arm, pos, nrm, view_pos, light_pos, min_roughness, False)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        kd, arm, pos, nrm, view_pos, light_pos = ctx.saved_variables
+        return _plugin.pbr_bsdf_bwd(kd, arm, pos, nrm, view_pos, light_pos, ctx.min_roughness, dout) + (None, None)
+
+def pbr_bsdf(kd, arm, pos, nrm, view_pos, light_pos, min_roughness=0.08, use_python=False):
+    '''Physically-based bsdf, both diffuse & specular lobes
+    All tensors assume a shape of [minibatch_size, height, width, 3] or broadcastable equivalent unless otherwise noted.
+
+    Args:
+        kd: Diffuse albedo.
+        arm: Specular parameters (attenuation, linear roughness, metalness).
+        pos: World space position.
+        nrm: World space shading normal.
+        view_pos: Camera position in world space, typically using broadcasting.
+        light_pos: Light position in world space, typically using broadcasting.
+        min_roughness: Scalar roughness clamping threshold
+
+        use_python: Use PyTorch implementation (for validation)
+
+    Returns:
+        Shaded color.
+    '''    
+
+    if use_python:
+        out = bsdf_pbr(kd, arm, pos, nrm, view_pos, light_pos, min_roughness=min_roughness)
+    else:
+        out = _pbr_bsdf_func.apply(kd, arm, pos, nrm, view_pos, light_pos, min_roughness)
+    
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of pbr_bsdf contains inf or NaN"
+    return out
+
+#----------------------------------------------------------------------------
+# Fast image loss function
+
+class _image_loss_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, img, target, loss, tonemapper):
+        ctx.loss, ctx.tonemapper = loss, tonemapper
+        ctx.save_for_backward(img, target)
+        out = _plugin.image_loss_fwd(img, target, loss, tonemapper, False)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        img, target = ctx.saved_variables
+        return _plugin.image_loss_bwd(img, target, dout, ctx.loss, ctx.tonemapper) + (None, None, None)
+
+def image_loss(img, target, loss='l1', tonemapper='none', use_python=False):
+    '''Compute HDR image loss. Combines tonemapping and loss into a single kernel for better perf.
+    All tensors assume a shape of [minibatch_size, height, width, 3] or broadcastable equivalent unless otherwise noted.
+
+    Args:
+        img: Input image.
+        target: Target (reference) image. 
+        loss: Type of loss. Valid options are ['l1', 'mse', 'smape', 'relmse']
+        tonemapper: Tonemapping operations. Valid options are ['none', 'log_srgb']
+        use_python: Use PyTorch implementation (for validation)
+
+    Returns:
+        Image space loss (scalar value).
+    '''
+    if use_python:
+        out = image_loss_fn(img, target, loss, tonemapper)
+    else:
+        out = _image_loss_func.apply(img, target, loss, tonemapper)
+        out = torch.sum(out) / (img.shape[0]*img.shape[1]*img.shape[2])
+
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of image_loss contains inf or NaN"
+    return out
+
+#----------------------------------------------------------------------------
+# Transform points function
+
+class _xfm_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, points, matrix, isPoints):
+        ctx.save_for_backward(points, matrix)
+        ctx.isPoints = isPoints
+        out = _plugin.xfm_fwd(points, matrix, isPoints, False)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        points, matrix = ctx.saved_variables
+        return (_plugin.xfm_bwd(points, matrix, dout, ctx.isPoints),) + (None, None, None)
+
+def xfm_points(points, matrix, use_python=False):
+    '''Transform points. 
+    Note: this method does not back-propagate matrix gradients by default for performance reasons. For matrix gradients, 
+    enable use_python=True or use torch.matmul instead.
+
+    Args:
+        points: Tensor containing 3D points with shape [minibatch_size, num_vertices, 3] or [1, num_vertices, 3]
+        matrix: A 4x4 transform matrix with shape [minibatch_size, 4, 4]
+        use_python: Use PyTorch's torch.matmul (for validation)
+    Returns:
+        Transformed points in homogeneous 4D with shape [minibatch_size, num_vertices, 4].
+    '''    
+    if use_python:
+        out = torch.matmul(torch.nn.functional.pad(points, pad=(0,1), mode='constant', value=1.0), torch.transpose(matrix, 1, 2))
+    else:
+        out = _xfm_func.apply(points, matrix, True)
+
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of xfm_points contains inf or NaN"
+    return out
+
+def xfm_vectors(vectors, matrix, use_python=False):
+    '''Transform vectors. 
+    Note: this method does not back-propagate matrix gradients by default for performance reasons. For matrix gradients, 
+    enable use_python=True or use torch.matmul instead.
+
+    Args:
+        vectors: Tensor containing 3D vectors with shape [minibatch_size, num_vertices, 3] or [1, num_vertices, 3]
+        matrix: A 4x4 transform matrix with shape [minibatch_size, 4, 4]
+        use_python: Use PyTorch's torch.matmul (for validation)
+
+    Returns:
+        Transformed vectors in homogeneous 4D with shape [minibatch_size, num_vertices, 4].
+    '''    
+
+    if use_python:
+        out = torch.matmul(torch.nn.functional.pad(vectors, pad=(0,1), mode='constant', value=0.0), torch.transpose(matrix, 1, 2))[..., 0:3].contiguous()
+    else:
+        out = _xfm_func.apply(vectors, matrix, False)
+
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of xfm_vectors contains inf or NaN"
+    return out

nvdiffmodeling/src/renderutils/tests/test_bsdf.py

@@ -0,0 +1,266 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import numpy as np
+import torch
+
+import os
+import sys
+sys.path.insert(0, os.path.join(sys.path[0], '../..'))
+import renderutils as ru
+
+RES = 4
+DTYPE = torch.float32
+
+def relative_loss(name, ref, cuda):
+	ref = ref.float()
+	cuda = cuda.float()
+	print(name, torch.max(torch.abs(ref - cuda) / torch.abs(ref + 1e-7)).item())
+
+def test_normal():
+	pos_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	pos_ref = pos_cuda.clone().detach().requires_grad_(True)
+	view_pos_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	view_pos_ref = view_pos_cuda.clone().detach().requires_grad_(True)
+	perturbed_nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	perturbed_nrm_ref = perturbed_nrm_cuda.clone().detach().requires_grad_(True)
+	smooth_nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	smooth_nrm_ref = smooth_nrm_cuda.clone().detach().requires_grad_(True)
+	smooth_tng_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	smooth_tng_ref = smooth_tng_cuda.clone().detach().requires_grad_(True)
+	geom_nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	geom_nrm_ref = geom_nrm_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda')
+
+	ref = ru.prepare_shading_normal(pos_ref, view_pos_ref, perturbed_nrm_ref, smooth_nrm_ref, smooth_tng_ref, geom_nrm_ref, True, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref, target)
+	ref_loss.backward()
+
+	cuda = ru.prepare_shading_normal(pos_cuda, view_pos_cuda, perturbed_nrm_cuda, smooth_nrm_cuda, smooth_tng_cuda, geom_nrm_cuda, True)
+	cuda_loss = torch.nn.MSELoss()(cuda, target)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+	print("    bent normal")
+	print("-------------------------------------------------------------")
+	relative_loss("res:", ref, cuda)
+	relative_loss("pos:", pos_ref.grad, pos_cuda.grad)
+	relative_loss("view_pos:", view_pos_ref.grad, view_pos_cuda.grad)
+	relative_loss("perturbed_nrm:", perturbed_nrm_ref.grad, perturbed_nrm_cuda.grad)
+	relative_loss("smooth_nrm:", smooth_nrm_ref.grad, smooth_nrm_cuda.grad)
+	relative_loss("smooth_tng:", smooth_tng_ref.grad, smooth_tng_cuda.grad)
+	relative_loss("geom_nrm:", geom_nrm_ref.grad, geom_nrm_cuda.grad)
+
+def test_schlick():
+	f0_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	f0_ref = f0_cuda.clone().detach().requires_grad_(True)
+	f90_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	f90_ref = f90_cuda.clone().detach().requires_grad_(True)
+	cosT_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) * 2.0
+	cosT_cuda = cosT_cuda.clone().detach().requires_grad_(True)
+	cosT_ref = cosT_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda')
+
+	ref = ru._fresnel_shlick(f0_ref, f90_ref, cosT_ref, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref, target)
+	ref_loss.backward()
+
+	cuda = ru._fresnel_shlick(f0_cuda, f90_cuda, cosT_cuda)
+	cuda_loss = torch.nn.MSELoss()(cuda, target)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+	print("    Fresnel shlick")
+	print("-------------------------------------------------------------")
+	relative_loss("res:", ref, cuda)
+	relative_loss("f0:", f0_ref.grad, f0_cuda.grad)
+	relative_loss("f90:", f90_ref.grad, f90_cuda.grad)
+	relative_loss("cosT:", cosT_ref.grad, cosT_cuda.grad)
+
+def test_ndf_ggx():
+	alphaSqr_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True)
+	alphaSqr_cuda = alphaSqr_cuda.clone().detach().requires_grad_(True)
+	alphaSqr_ref = alphaSqr_cuda.clone().detach().requires_grad_(True)
+	cosT_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) * 3.0 - 1
+	cosT_cuda = cosT_cuda.clone().detach().requires_grad_(True)
+	cosT_ref = cosT_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda')
+
+	ref = ru._ndf_ggx(alphaSqr_ref, cosT_ref, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref, target)
+	ref_loss.backward()
+
+	cuda = ru._ndf_ggx(alphaSqr_cuda, cosT_cuda)
+	cuda_loss = torch.nn.MSELoss()(cuda, target)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+	print("    Ndf GGX")
+	print("-------------------------------------------------------------")
+	relative_loss("res:", ref, cuda)
+	relative_loss("alpha:", alphaSqr_ref.grad, alphaSqr_cuda.grad)
+	relative_loss("cosT:", cosT_ref.grad, cosT_cuda.grad)
+
+def test_lambda_ggx():
+	alphaSqr_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True)
+	alphaSqr_ref = alphaSqr_cuda.clone().detach().requires_grad_(True)
+	cosT_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) * 3.0 - 1
+	cosT_cuda = cosT_cuda.clone().detach().requires_grad_(True)
+	cosT_ref = cosT_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda')
+
+	ref = ru._lambda_ggx(alphaSqr_ref, cosT_ref, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref, target)
+	ref_loss.backward()
+
+	cuda = ru._lambda_ggx(alphaSqr_cuda, cosT_cuda)
+	cuda_loss = torch.nn.MSELoss()(cuda, target)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+	print("    Lambda GGX")
+	print("-------------------------------------------------------------")
+	relative_loss("res:", ref, cuda)
+	relative_loss("alpha:", alphaSqr_ref.grad, alphaSqr_cuda.grad)
+	relative_loss("cosT:", cosT_ref.grad, cosT_cuda.grad)
+
+def test_masking_smith():
+	alphaSqr_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True)
+	alphaSqr_ref = alphaSqr_cuda.clone().detach().requires_grad_(True)
+	cosThetaI_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True)
+	cosThetaI_ref = cosThetaI_cuda.clone().detach().requires_grad_(True)
+	cosThetaO_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True)
+	cosThetaO_ref = cosThetaO_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda')
+
+	ref = ru._masking_smith(alphaSqr_ref, cosThetaI_ref, cosThetaO_ref, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref, target)
+	ref_loss.backward()
+
+	cuda = ru._masking_smith(alphaSqr_cuda, cosThetaI_cuda, cosThetaO_cuda)
+	cuda_loss = torch.nn.MSELoss()(cuda, target)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+	print("    Smith masking term")
+	print("-------------------------------------------------------------")
+	relative_loss("res:", ref, cuda)
+	relative_loss("alpha:", alphaSqr_ref.grad, alphaSqr_cuda.grad)
+	relative_loss("cosThetaI:", cosThetaI_ref.grad, cosThetaI_cuda.grad)
+	relative_loss("cosThetaO:", cosThetaO_ref.grad, cosThetaO_cuda.grad)
+
+def test_lambert():
+	normals_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	normals_ref = normals_cuda.clone().detach().requires_grad_(True)
+	wi_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	wi_ref = wi_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda')
+
+	ref = ru.lambert(normals_ref, wi_ref, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref, target)
+	ref_loss.backward()
+
+	cuda = ru.lambert(normals_cuda, wi_cuda)
+	cuda_loss = torch.nn.MSELoss()(cuda, target)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+	print("    Lambert")
+	print("-------------------------------------------------------------")
+	relative_loss("res:", ref, cuda)
+	relative_loss("nrm:", normals_ref.grad, normals_cuda.grad)
+	relative_loss("wi:", wi_ref.grad, wi_cuda.grad)
+
+def test_pbr_specular():
+	col_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	col_ref = col_cuda.clone().detach().requires_grad_(True)
+	nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	nrm_ref = nrm_cuda.clone().detach().requires_grad_(True)
+	wi_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	wi_ref = wi_cuda.clone().detach().requires_grad_(True)
+	wo_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	wo_ref = wo_cuda.clone().detach().requires_grad_(True)
+	alpha_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True)
+	alpha_ref = alpha_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda')
+
+	ref = ru.pbr_specular(col_ref, nrm_ref, wo_ref, wi_ref, alpha_ref, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref, target)
+	ref_loss.backward()
+
+	cuda = ru.pbr_specular(col_cuda, nrm_cuda, wo_cuda, wi_cuda, alpha_cuda)
+	cuda_loss = torch.nn.MSELoss()(cuda, target)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+	print("    Pbr specular")
+	print("-------------------------------------------------------------")
+
+	relative_loss("res:", ref, cuda)
+	if col_ref.grad is not None:
+		relative_loss("col:", col_ref.grad, col_cuda.grad)
+	if nrm_ref.grad is not None:
+		relative_loss("nrm:", nrm_ref.grad, nrm_cuda.grad)
+	if wi_ref.grad is not None:
+		relative_loss("wi:", wi_ref.grad, wi_cuda.grad)
+	if wo_ref.grad is not None:
+		relative_loss("wo:", wo_ref.grad, wo_cuda.grad)
+	if alpha_ref.grad is not None:
+		relative_loss("alpha:", alpha_ref.grad, alpha_cuda.grad)
+
+def test_pbr_bsdf():
+	kd_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	kd_ref = kd_cuda.clone().detach().requires_grad_(True)
+	arm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	arm_ref = arm_cuda.clone().detach().requires_grad_(True)
+	pos_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	pos_ref = pos_cuda.clone().detach().requires_grad_(True)
+	nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	nrm_ref = nrm_cuda.clone().detach().requires_grad_(True)
+	view_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	view_ref = view_cuda.clone().detach().requires_grad_(True)
+	light_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	light_ref = light_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda')
+
+	ref = ru.pbr_bsdf(kd_ref, arm_ref, pos_ref, nrm_ref, view_ref, light_ref, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref, target)
+	ref_loss.backward()
+
+	cuda = ru.pbr_bsdf(kd_cuda, arm_cuda, pos_cuda, nrm_cuda, view_cuda, light_cuda)
+	cuda_loss = torch.nn.MSELoss()(cuda, target)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+	print("    Pbr BSDF")
+	print("-------------------------------------------------------------")
+
+	relative_loss("res:", ref, cuda)
+	if kd_ref.grad is not None:
+		relative_loss("kd:", kd_ref.grad, kd_cuda.grad)
+	if arm_ref.grad is not None:
+		relative_loss("arm:", arm_ref.grad, arm_cuda.grad)
+	if pos_ref.grad is not None:
+		relative_loss("pos:", pos_ref.grad, pos_cuda.grad)
+	if nrm_ref.grad is not None:
+		relative_loss("nrm:", nrm_ref.grad, nrm_cuda.grad)
+	if view_ref.grad is not None:
+		relative_loss("view:", view_ref.grad, view_cuda.grad)
+	if light_ref.grad is not None:
+		relative_loss("light:", light_ref.grad, light_cuda.grad)
+
+test_normal()
+
+test_schlick()
+test_ndf_ggx()
+test_lambda_ggx()
+test_masking_smith()
+
+test_lambert()
+test_pbr_specular()
+test_pbr_bsdf()

nvdiffmodeling/src/renderutils/tests/test_loss.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import numpy as np
+import torch
+
+import os
+import sys
+sys.path.insert(0, os.path.join(sys.path[0], '../..'))
+import renderutils as ru
+
+RES = 8
+DTYPE = torch.float32
+
+def tonemap_srgb(f):
+    return torch.where(f > 0.0031308, torch.pow(torch.clamp(f, min=0.0031308), 1.0/2.4)*1.055 - 0.055, 12.92*f)
+
+def l1(output, target):
+    x = torch.clamp(output, min=0, max=65535)
+    r = torch.clamp(target, min=0, max=65535)
+    x = tonemap_srgb(torch.log(x + 1))
+    r = tonemap_srgb(torch.log(r + 1))
+    return torch.nn.functional.l1_loss(x,r)
+
+def relative_loss(name, ref, cuda):
+	ref = ref.float()
+	cuda = cuda.float()
+	print(name, torch.max(torch.abs(ref - cuda) / torch.abs(ref + 1e-7)).item())
+
+def test_loss(loss, tonemapper):
+	img_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	img_ref = img_cuda.clone().detach().requires_grad_(True)
+	target_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	target_ref = target_cuda.clone().detach().requires_grad_(True)
+
+	ref_loss = ru.image_loss(img_ref, target_ref, loss=loss, tonemapper=tonemapper, use_python=True)
+	ref_loss.backward()
+
+	cuda_loss = ru.image_loss(img_cuda, target_cuda, loss=loss, tonemapper=tonemapper)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+	print("    Loss: %s, %s" % (loss, tonemapper))
+	print("-------------------------------------------------------------")
+
+	relative_loss("res:", ref_loss, cuda_loss)
+	relative_loss("img:", img_ref.grad, img_cuda.grad)
+	relative_loss("target:", target_ref.grad, target_cuda.grad)
+
+
+test_loss('l1', 'none')
+test_loss('l1', 'log_srgb')
+test_loss('mse', 'log_srgb')
+test_loss('smape', 'none')
+test_loss('relmse', 'none')
+test_loss('mse', 'none')

nvdiffmodeling/src/renderutils/tests/test_mesh.py

@@ -0,0 +1,90 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import numpy as np
+import torch
+
+import os
+import sys
+sys.path.insert(0, os.path.join(sys.path[0], '../..'))
+import renderutils as ru
+
+BATCH = 8
+RES = 1024
+DTYPE = torch.float32
+
+torch.manual_seed(0)
+
+def tonemap_srgb(f):
+    return torch.where(f > 0.0031308, torch.pow(torch.clamp(f, min=0.0031308), 1.0/2.4)*1.055 - 0.055, 12.92*f)
+
+def l1(output, target):
+    x = torch.clamp(output, min=0, max=65535)
+    r = torch.clamp(target, min=0, max=65535)
+    x = tonemap_srgb(torch.log(x + 1))
+    r = tonemap_srgb(torch.log(r + 1))
+    return torch.nn.functional.l1_loss(x,r)
+
+def relative_loss(name, ref, cuda):
+	ref = ref.float()
+	cuda = cuda.float()
+	print(name, torch.max(torch.abs(ref - cuda) / torch.abs(ref)).item())
+
+def test_xfm_points():
+	points_cuda = torch.rand(1, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	points_ref = points_cuda.clone().detach().requires_grad_(True)
+	mtx_cuda = torch.rand(BATCH, 4, 4, dtype=DTYPE, device='cuda', requires_grad=False)
+	mtx_ref = mtx_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(BATCH, RES, 4, dtype=DTYPE, device='cuda', requires_grad=True)
+
+	ref_out = ru.xfm_points(points_ref, mtx_ref, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref_out, target)
+	ref_loss.backward()
+
+	cuda_out = ru.xfm_points(points_cuda, mtx_cuda)
+	cuda_loss = torch.nn.MSELoss()(cuda_out, target)
+	cuda_loss.backward()
+
+	print("-------------------------------------------------------------")
+
+	relative_loss("res:", ref_out, cuda_out)
+	relative_loss("points:", points_ref.grad, points_cuda.grad)
+
+def test_xfm_vectors():
+	points_cuda = torch.rand(1, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	points_ref = points_cuda.clone().detach().requires_grad_(True)
+	points_cuda_p = points_cuda.clone().detach().requires_grad_(True)
+	points_ref_p = points_cuda.clone().detach().requires_grad_(True)
+	mtx_cuda = torch.rand(BATCH, 4, 4, dtype=DTYPE, device='cuda', requires_grad=False)
+	mtx_ref = mtx_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(BATCH, RES, 4, dtype=DTYPE, device='cuda', requires_grad=True)
+
+	ref_out = ru.xfm_vectors(points_ref.contiguous(), mtx_ref, use_python=True)
+	ref_loss = torch.nn.MSELoss()(ref_out, target[..., 0:3])
+	ref_loss.backward()
+
+	cuda_out = ru.xfm_vectors(points_cuda.contiguous(), mtx_cuda)
+	cuda_loss = torch.nn.MSELoss()(cuda_out, target[..., 0:3])
+	cuda_loss.backward()
+
+	ref_out_p = ru.xfm_points(points_ref_p.contiguous(), mtx_ref, use_python=True)
+	ref_loss_p = torch.nn.MSELoss()(ref_out_p, target)
+	ref_loss_p.backward()
+	
+	cuda_out_p = ru.xfm_points(points_cuda_p.contiguous(), mtx_cuda)
+	cuda_loss_p = torch.nn.MSELoss()(cuda_out_p, target)
+	cuda_loss_p.backward()
+
+	print("-------------------------------------------------------------")
+
+	relative_loss("res:", ref_out, cuda_out)
+	relative_loss("points:", points_ref.grad, points_cuda.grad)
+	relative_loss("points_p:", points_ref_p.grad, points_cuda_p.grad)
+
+test_xfm_points()
+test_xfm_vectors()

nvdiffmodeling/src/renderutils/tests/test_perf.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import numpy as np
+import torch
+
+import os
+import sys
+import time
+sys.path.insert(0, os.path.join(sys.path[0], '../..'))
+import renderutils as ru
+
+DTYPE=torch.float32
+
+def test_bsdf(BATCH, RES, ITR):
+	kd_cuda = torch.rand(BATCH, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	kd_ref = kd_cuda.clone().detach().requires_grad_(True)
+	arm_cuda = torch.rand(BATCH, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	arm_ref = arm_cuda.clone().detach().requires_grad_(True)
+	pos_cuda = torch.rand(BATCH, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	pos_ref = pos_cuda.clone().detach().requires_grad_(True)
+	nrm_cuda = torch.rand(BATCH, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	nrm_ref = nrm_cuda.clone().detach().requires_grad_(True)
+	view_cuda = torch.rand(BATCH, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	view_ref = view_cuda.clone().detach().requires_grad_(True)
+	light_cuda = torch.rand(BATCH, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True)
+	light_ref = light_cuda.clone().detach().requires_grad_(True)
+	target = torch.rand(BATCH, RES, RES, 3, device='cuda')
+
+	start = torch.cuda.Event(enable_timing=True)
+	end = torch.cuda.Event(enable_timing=True)
+
+	ru.pbr_bsdf(kd_cuda, arm_cuda, pos_cuda, nrm_cuda, view_cuda, light_cuda)
+
+	print("--- Testing: [%d, %d, %d] ---" % (BATCH, RES, RES))
+
+	start.record()
+	for i in range(ITR):
+		ref = ru.pbr_bsdf(kd_ref, arm_ref, pos_ref, nrm_ref, view_ref, light_ref, use_python=True)
+	end.record()
+	torch.cuda.synchronize()
+	print("Pbr BSDF python:", start.elapsed_time(end))
+
+	start.record()
+	for i in range(ITR):
+		cuda = ru.pbr_bsdf(kd_cuda, arm_cuda, pos_cuda, nrm_cuda, view_cuda, light_cuda)
+	end.record()
+	torch.cuda.synchronize()
+	print("Pbr BSDF cuda:", start.elapsed_time(end))
+
+test_bsdf(1, 512, 1000)
+test_bsdf(16, 512, 1000)
+test_bsdf(1, 2048, 1000)

nvdiffmodeling/src/texture.py

@@ -0,0 +1,151 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import numpy as np
+import torch
+import nvdiffrast.torch as dr
+
+from . import util
+
+########################################################################################################
+# Simple texture class. A texture can be either 
+# - A 3D tensor (using auto mipmaps)
+# - A list of 3D tensors (full custom mip hierarchy)
+########################################################################################################
+
+class Texture2D:
+     # Initializes a texture from image data.
+     # Input can be constant value (1D array) or texture (3D array) or mip hierarchy (list of 3d arrays)
+    def __init__(self, init):
+        if isinstance(init, np.ndarray):
+            init = torch.tensor(init, dtype=torch.float32, device='cuda')
+        elif isinstance(init, list) and len(init) == 1:
+            init = init[0]
+
+        if isinstance(init, list) or len(init.shape) == 4:
+            self.data = init
+        elif len(init.shape) == 3:
+            self.data = init[None, ...]
+        else:
+            self.data = init[None, None, None, :] # Convert constant to 1x1 tensor
+
+    # Filtered (trilinear) sample texture at a given location
+    def sample(self, texc, texc_deriv, filter_mode='linear-mipmap-linear', data_fmt=torch.float32):
+        if isinstance(self.data, list):
+            out = dr.texture(self.data[0], texc, texc_deriv, mip=self.data[1:], filter_mode=filter_mode)
+        else:
+            out = dr.texture(self.data, texc, texc_deriv, filter_mode=filter_mode)
+        return out.to(data_fmt)
+
+    def getRes(self):
+        return self.getMips()[0].shape[1:3]
+
+    def getMips(self):
+        if isinstance(self.data, list):
+            return self.data
+        else:
+            return [self.data]
+
+    # In-place clamp with no derivative to make sure values are in valid range after training
+    def clamp_(self, min=None, max=None):
+        with torch.no_grad():
+            for mip in self.getMips():
+                mip.clamp_(min=min, max=max)
+
+    # In-place clamp with no derivative to make sure values are in valid range after training
+    def clamp_rgb_(self, minR=None, maxR=None, minG=None, maxG=None, minB=None, maxB=None):
+        with torch.no_grad():
+            for mip in self.getMips():
+                mip[...,0].clamp_(min=minR, max=maxR)
+                mip[...,1].clamp_(min=minG, max=maxG)
+                mip[...,2].clamp_(min=minB, max=maxB)
+
+########################################################################################################
+# Helper function to create a trainable texture from a regular texture. The trainable weights are 
+# initialized with texture data as an initial guess
+########################################################################################################
+
+def create_trainable(init, res, auto_mipmaps):
+    with torch.no_grad():
+        if isinstance(init, Texture2D):
+            assert isinstance(init.data, torch.Tensor)
+            init = init.data
+        elif isinstance(init, np.ndarray):
+            init = torch.tensor(init, dtype=torch.float32, device='cuda')
+
+        # Pad to NHWC if needed
+        if len(init.shape) == 1: # Extend constant to NHWC tensor
+            init = init[None, None, None, :]
+        elif len(init.shape) == 3:
+            init = init[None, ...]
+
+        # Scale input to desired resolution.
+        init = util.scale_img_nhwc(init, res)
+
+        # Genreate custom mipchain
+        if not auto_mipmaps:
+            mip_chain = [init.clone().detach().requires_grad_(True)]
+            while mip_chain[-1].shape[1] > 1 or mip_chain[-1].shape[2] > 1:
+                new_size = [max(mip_chain[-1].shape[1] // 2, 1), max(mip_chain[-1].shape[2] // 2, 1)]
+                init = util.scale_img_nhwc(mip_chain[-1], new_size)
+                mip_chain += [init.clone().detach().requires_grad_(True)]
+            return Texture2D(mip_chain)
+        else:
+            return Texture2D(init.clone().detach().requires_grad_(True))
+
+########################################################################################################
+# Convert texture to and from SRGB
+########################################################################################################
+
+def srgb_to_rgb(texture):
+    return Texture2D(list(util.srgb_to_rgb(mip) for mip in texture.getMips()))
+
+def rgb_to_srgb(texture):
+    return Texture2D(list(util.rgb_to_srgb(mip) for mip in texture.getMips()))
+
+########################################################################################################
+# Utility functions for loading / storing a texture
+########################################################################################################
+
+def _load_mip2D(fn, lambda_fn=None, channels=None):
+    imgdata = torch.tensor(util.load_image(fn), dtype=torch.float32, device='cuda')
+    if channels is not None:
+        imgdata = imgdata[..., 0:channels]
+    if lambda_fn is not None:
+        imgdata = lambda_fn(imgdata)
+    return imgdata.detach().clone()
+
+def load_texture2D(fn, lambda_fn=None, channels=None):
+    base, ext = os.path.splitext(fn)
+    if os.path.exists(base + "_0" + ext):
+        mips = []
+        while os.path.exists(base + ("_%d" % len(mips)) + ext):
+            mips += [_load_mip2D(base + ("_%d" % len(mips)) + ext, lambda_fn, channels)]
+        return Texture2D(mips)
+    else:
+        return Texture2D(_load_mip2D(fn, lambda_fn, channels))
+
+def _save_mip2D(fn, mip, mipidx, lambda_fn):
+    if lambda_fn is not None:
+        data = lambda_fn(mip).detach().cpu().numpy()
+    else:
+        data = mip.detach().cpu().numpy()
+
+    if mipidx is None:
+        util.save_image(fn, data)
+    else:
+        base, ext = os.path.splitext(fn)
+        util.save_image(base + ("_%d" % mipidx) + ext, data)
+
+def save_texture2D(fn, tex, lambda_fn=None):
+    if isinstance(tex.data, list):
+        for i, mip in enumerate(tex.data):
+            _save_mip2D(fn, mip[0,...], i, lambda_fn)
+    else:
+        _save_mip2D(fn, tex.data[0,...], None, lambda_fn)

nvdiffmodeling/src/util.py

@@ -0,0 +1,354 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto. Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import sys
+import numpy as np
+import torch
+import nvdiffrast.torch as dr
+import imageio
+
+#----------------------------------------------------------------------------
+# Vector operations
+#----------------------------------------------------------------------------
+
+def dot(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    return torch.sum(x*y, -1, keepdim=True)
+
+def reflect(x: torch.Tensor, n: torch.Tensor) -> torch.Tensor:
+    return 2*dot(x, n)*n - x
+
+def length(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor:
+    return torch.sqrt(torch.clamp(dot(x,x), min=eps)) # Clamp to avoid nan gradients because grad(sqrt(0)) = NaN
+
+def safe_normalize(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor:
+    return x / length(x, eps)
+
+def to_hvec(x: torch.Tensor, w: float) -> torch.Tensor:
+    return torch.nn.functional.pad(x, pad=(0,1), mode='constant', value=w)
+
+#----------------------------------------------------------------------------
+# Tonemapping
+#----------------------------------------------------------------------------
+
+def tonemap_srgb(f: torch.Tensor) -> torch.Tensor:
+    return torch.where(f > 0.0031308, torch.pow(torch.clamp(f, min=0.0031308), 1.0/2.4)*1.055 - 0.055, 12.92*f)
+
+#----------------------------------------------------------------------------
+# sRGB color transforms
+#----------------------------------------------------------------------------
+
+def _rgb_to_srgb(f: torch.Tensor) -> torch.Tensor:
+    return torch.where(f <= 0.0031308, f * 12.92, torch.pow(torch.clamp(f, 0.0031308), 1.0/2.4)*1.055 - 0.055)
+
+def rgb_to_srgb(f: torch.Tensor) -> torch.Tensor:
+    assert f.shape[-1] == 3 or f.shape[-1] == 4
+    out = torch.cat((_rgb_to_srgb(f[..., 0:3]), f[..., 3:4]), dim=-1) if f.shape[-1] == 4 else _rgb_to_srgb(f)
+    assert out.shape[0] == f.shape[0] and out.shape[1] == f.shape[1] and out.shape[2] == f.shape[2]
+    return out
+
+def _srgb_to_rgb(f: torch.Tensor) -> torch.Tensor:
+    return torch.where(f <= 0.04045, f / 12.92, torch.pow((torch.clamp(f, 0.04045) + 0.055) / 1.055, 2.4))
+
+def srgb_to_rgb(f: torch.Tensor) -> torch.Tensor:
+    assert f.shape[-1] == 3 or f.shape[-1] == 4
+    out = torch.cat((_srgb_to_rgb(f[..., 0:3]), f[..., 3:4]), dim=-1) if f.shape[-1] == 4 else _srgb_to_rgb(f)
+    assert out.shape[0] == f.shape[0] and out.shape[1] == f.shape[1] and out.shape[2] == f.shape[2]
+    return out
+
+#----------------------------------------------------------------------------
+# Displacement texture lookup
+#----------------------------------------------------------------------------
+
+def get_miplevels(texture: np.ndarray) -> float:
+    minDim = min(texture.shape[0], texture.shape[1])
+    return np.floor(np.log2(minDim))
+
+# TODO: Handle wrapping maybe
+def tex_2d(tex_map : torch.Tensor, coords : torch.Tensor, filter='nearest') -> torch.Tensor:
+    tex_map = tex_map[None, ...]    # Add batch dimension
+    tex_map = tex_map.permute(0, 3, 1, 2) # NHWC -> NCHW
+    tex = torch.nn.functional.grid_sample(tex_map, coords[None, None, ...] * 2 - 1, mode=filter, align_corners=False)
+    tex = tex.permute(0, 2, 3, 1) # NCHW -> NHWC
+    return tex[0, 0, ...]
+
+#----------------------------------------------------------------------------
+# Image scaling
+#----------------------------------------------------------------------------
+
+def scale_img_hwc(x : torch.Tensor, size, mag='bilinear', min='area') -> torch.Tensor:
+    return scale_img_nhwc(x[None, ...], size, mag, min)[0]
+
+def scale_img_nhwc(x  : torch.Tensor, size, mag='bilinear', min='area') -> torch.Tensor:
+    assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] < size[0] and x.shape[2] < size[1]), "Trying to magnify image in one dimension and minify in the other"
+    y = x.permute(0, 3, 1, 2) # NHWC -> NCHW
+    if x.shape[1] > size[0] and x.shape[2] > size[1]: # Minification, previous size was bigger
+        y = torch.nn.functional.interpolate(y, size, mode=min)
+    else: # Magnification
+        if mag == 'bilinear' or mag == 'bicubic':
+            y = torch.nn.functional.interpolate(y, size, mode=mag, align_corners=True)
+        else:
+            y = torch.nn.functional.interpolate(y, size, mode=mag)
+    return y.permute(0, 2, 3, 1).contiguous() # NCHW -> NHWC
+
+def avg_pool_nhwc(x  : torch.Tensor, size) -> torch.Tensor:
+    y = x.permute(0, 3, 1, 2) # NHWC -> NCHW
+    y = torch.nn.functional.avg_pool2d(y, size)
+    return y.permute(0, 2, 3, 1).contiguous() # NCHW -> NHWC
+
+#----------------------------------------------------------------------------
+# Behaves similar to tf.segment_sum
+#----------------------------------------------------------------------------
+
+def segment_sum(data: torch.Tensor, segment_ids: torch.Tensor) -> torch.Tensor:
+    num_segments = torch.unique_consecutive(segment_ids).shape[0]
+
+    # Repeats ids until same dimension as data
+    if len(segment_ids.shape) == 1:
+        s = torch.prod(torch.tensor(data.shape[1:], dtype=torch.int64, device='cuda')).long()
+        segment_ids = segment_ids.repeat_interleave(s).view(segment_ids.shape[0], *data.shape[1:])
+
+    assert data.shape == segment_ids.shape, "data.shape and segment_ids.shape should be equal"
+
+    shape = [num_segments] + list(data.shape[1:])
+    result = torch.zeros(*shape, dtype=torch.float32, device='cuda')
+    result = result.scatter_add(0, segment_ids, data)
+    return result
+
+#----------------------------------------------------------------------------
+# Projection and transformation matrix helpers.
+#----------------------------------------------------------------------------
+
+def projection(x=0.1, n=1.0, f=50.0):
+    return np.array([[n/x,    0,            0,              0], 
+                     [  0, n/-x,            0,              0], 
+                     [  0,    0, -(f+n)/(f-n), -(2*f*n)/(f-n)], 
+                     [  0,    0,           -1,              0]]).astype(np.float32)
+                    
+def translate(x, y, z):
+    return np.array([[1, 0, 0, x], 
+                     [0, 1, 0, y], 
+                     [0, 0, 1, z], 
+                     [0, 0, 0, 1]]).astype(np.float32)
+
+def rotate_x(a):
+    s, c = np.sin(a), np.cos(a)
+    return np.array([[1,  0, 0, 0], 
+                     [0,  c, s, 0], 
+                     [0, -s, c, 0], 
+                     [0,  0, 0, 1]]).astype(np.float32)
+
+def rotate_y(a):
+    s, c = np.sin(a), np.cos(a)
+    return np.array([[ c, 0, s, 0], 
+                     [ 0, 1, 0, 0], 
+                     [-s, 0, c, 0], 
+                     [ 0, 0, 0, 1]]).astype(np.float32)
+
+def scale(s):
+    return np.array([[ s, 0, 0, 0], 
+                     [ 0, s, 0, 0], 
+                     [ 0, 0, s, 0], 
+                     [ 0, 0, 0, 1]]).astype(np.float32)
+
+def lookAt(eye, at, up):
+    a = eye - at
+    b = up
+    w = a / np.linalg.norm(a)
+    u = np.cross(b, w)
+    u = u / np.linalg.norm(u)
+    v = np.cross(w, u)
+    translate = np.array([[1, 0, 0, -eye[0]], 
+                          [0, 1, 0, -eye[1]], 
+                          [0, 0, 1, -eye[2]], 
+                          [0, 0, 0, 1]]).astype(np.float32)
+    rotate =  np.array([[u[0], u[1], u[2], 0], 
+                        [v[0], v[1], v[2], 0], 
+                        [w[0], w[1], w[2], 0], 
+                        [0, 0, 0, 1]]).astype(np.float32)
+    return np.matmul(rotate, translate)
+
+def random_rotation_translation(t):
+    m = np.random.normal(size=[3, 3])
+    m[1] = np.cross(m[0], m[2])
+    m[2] = np.cross(m[0], m[1])
+    m = m / np.linalg.norm(m, axis=1, keepdims=True)
+    m = np.pad(m, [[0, 1], [0, 1]], mode='constant')
+    m[3, 3] = 1.0
+    m[:3, 3] = np.random.uniform(-t, t, size=[3])
+    return m
+
+
+#----------------------------------------------------------------------------
+# Cosine sample around a vector N
+#----------------------------------------------------------------------------
+def cosine_sample(N : np.ndarray) -> np.ndarray:
+    # construct local frame
+    N = N/np.linalg.norm(N)
+
+    dx0 = np.array([0, N[2], -N[1]])
+    dx1 = np.array([-N[2], 0, N[0]])
+
+    dx = dx0 if np.dot(dx0,dx0) > np.dot(dx1,dx1) else dx1
+    dx = dx/np.linalg.norm(dx)
+    dy = np.cross(N,dx)
+    dy = dy/np.linalg.norm(dy)
+
+    # cosine sampling in local frame
+    phi = 2.0*np.pi*np.random.uniform()
+    s = np.random.uniform()
+    costheta = np.sqrt(s)
+    sintheta = np.sqrt(1.0 - s)
+
+    # cartesian vector in local space
+    x = np.cos(phi)*sintheta
+    y = np.sin(phi)*sintheta
+    z = costheta
+
+    # local to world
+    return dx*x + dy*y + N*z
+
+
+#----------------------------------------------------------------------------
+# Cosine sampled light directions around the vector N
+#----------------------------------------------------------------------------
+def cosine_sample_texture(res, N : np.ndarray) -> torch.Tensor:
+    # construct local frame
+    N = N/np.linalg.norm(N)
+
+    dx0 = np.array([0, N[2], -N[1]])
+    dx1 = np.array([-N[2], 0, N[0]])
+
+    dx = dx0 if np.dot(dx0,dx0) > np.dot(dx1,dx1) else dx1
+    dx = dx/np.linalg.norm(dx)
+    dy = np.cross(N,dx)
+    dy = dy/np.linalg.norm(dy)
+
+    X = torch.tensor(dx, dtype=torch.float32, device='cuda')
+    Y = torch.tensor(dy, dtype=torch.float32, device='cuda')
+    Z = torch.tensor(N, dtype=torch.float32, device='cuda')
+
+    # cosine sampling in local frame
+
+    phi = 2.0*np.pi*torch.rand(res, res, 1, dtype=torch.float32, device='cuda')
+    s = torch.rand(res, res, 1, dtype=torch.float32, device='cuda')
+    costheta = torch.sqrt(s)
+    sintheta = torch.sqrt(1.0 - s)
+
+    # cartesian vector in local space
+    x = torch.cos(phi)*sintheta
+    y = torch.sin(phi)*sintheta
+    z = costheta
+
+    # local to world
+    return X*x + Y*y + Z*z
+
+#----------------------------------------------------------------------------
+# Bilinear downsample by 2x.
+#----------------------------------------------------------------------------
+
+def bilinear_downsample(x : torch.tensor) -> torch.Tensor:
+    w = torch.tensor([[1, 3, 3, 1], [3, 9, 9, 3], [3, 9, 9, 3], [1, 3, 3, 1]], dtype=torch.float32, device=x.device) / 64.0
+    w = w.expand(x.shape[-1], 1, 4, 4) 
+    x = torch.nn.functional.conv2d(x.permute(0, 3, 1, 2), w, padding=1, stride=2, groups=x.shape[-1])
+    return x.permute(0, 2, 3, 1)
+
+#----------------------------------------------------------------------------
+# Bilinear downsample log(spp) steps
+#----------------------------------------------------------------------------
+
+def bilinear_downsample(x : torch.tensor, spp) -> torch.Tensor:
+    w = torch.tensor([[1, 3, 3, 1], [3, 9, 9, 3], [3, 9, 9, 3], [1, 3, 3, 1]], dtype=torch.float32, device=x.device) / 64.0
+    g = x.shape[-1]
+    w = w.expand(g, 1, 4, 4) 
+    x = x.permute(0, 3, 1, 2) # NHWC -> NCHW
+    steps = int(np.log2(spp))
+    for _ in range(steps):
+        xp = torch.nn.functional.pad(x, (1,1,1,1), mode='replicate')
+        x = torch.nn.functional.conv2d(xp, w, padding=0, stride=2, groups=g)
+    return x.permute(0, 2, 3, 1).contiguous() # NCHW -> NHWC
+
+
+#----------------------------------------------------------------------------
+# Image display function using OpenGL.
+#----------------------------------------------------------------------------
+
+_glfw_window = None
+def display_image(image, zoom=None, size=None, title=None): # HWC
+    # Import OpenGL and glfw.
+    import OpenGL.GL as gl
+    import glfw
+
+    # Zoom image if requested.
+    image = np.asarray(image)
+    if size is not None:
+        assert zoom is None
+        zoom = max(1, size // image.shape[0])
+    if zoom is not None:
+        image = image.repeat(zoom, axis=0).repeat(zoom, axis=1)
+    height, width, channels = image.shape
+
+    # Initialize window.
+    if title is None:
+        title = 'Debug window'
+    global _glfw_window
+    if _glfw_window is None:
+        glfw.init()
+        _glfw_window = glfw.create_window(width, height, title, None, None)
+        glfw.make_context_current(_glfw_window)
+        glfw.show_window(_glfw_window)
+        glfw.swap_interval(0)
+    else:
+        glfw.make_context_current(_glfw_window)
+        glfw.set_window_title(_glfw_window, title)
+        glfw.set_window_size(_glfw_window, width, height)
+
+    # Update window.
+    glfw.poll_events()
+    gl.glClearColor(0, 0, 0, 1)
+    gl.glClear(gl.GL_COLOR_BUFFER_BIT)
+    gl.glWindowPos2f(0, 0)
+    gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
+    gl_format = {3: gl.GL_RGB, 2: gl.GL_RG, 1: gl.GL_LUMINANCE}[channels]
+    gl_dtype = {'uint8': gl.GL_UNSIGNED_BYTE, 'float32': gl.GL_FLOAT}[image.dtype.name]
+    gl.glDrawPixels(width, height, gl_format, gl_dtype, image[::-1])
+    glfw.swap_buffers(_glfw_window)
+    if glfw.window_should_close(_glfw_window):
+        return False
+    return True
+
+#----------------------------------------------------------------------------
+# Image save helper.
+#----------------------------------------------------------------------------
+
+def save_image(fn, x : np.ndarray) -> np.ndarray:
+    imageio.imwrite(fn, np.clip(np.rint(x * 255.0), 0, 255).astype(np.uint8))
+
+def load_image(fn) -> np.ndarray:
+    img = imageio.imread(fn)
+    if img.dtype == np.float32: # HDR image
+        return img
+    else: # LDR image
+        return img.astype(np.float32) / 255
+
+#----------------------------------------------------------------------------
+
+def time_to_text(x):
+    if x > 3600:
+        return "%.2f h" % (x / 3600)
+    elif x > 60:
+        return "%.2f m" % (x / 60)
+    else:
+        return "%.2f s" % x
+
+#----------------------------------------------------------------------------
+
+def checkerboard(width, repetitions) -> np.ndarray:
+    tilesize = int(width//repetitions//2)
+    check = np.kron([[1, 0] * repetitions, [0, 1] * repetitions] * repetitions, np.ones((tilesize, tilesize)))*0.33 + 0.33
+    return np.stack((check, check, check), axis=-1)[None, ...]