Add spar3d demo files

.gitattributes

@@ -33,3 +33,5 @@ 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
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,167 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv*/
+env/
+venv*/
+ENV/
+env.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+.vs/
+.idea/
+.vscode/
+
+stabilityai/
+output/

.pre-commit-config.yaml

@@ -0,0 +1,24 @@
+default_language_version:
+  python: python3
+
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
+    hooks:
+      - id: trailing-whitespace
+      - id: check-ast
+      - id: check-merge-conflict
+      - id: check-yaml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+        args: [--markdown-linebreak-ext=md]
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.3.5
+    hooks:
+      # Run the linter.
+      - id: ruff
+        args: [ --fix ]
+      # Run the formatter.
+      - id: ruff-format

LICENSE.md

@@ -0,0 +1,51 @@
+STABILITY AI COMMUNITY LICENSE AGREEMENT
+Last Updated: July 5, 2024
+
+
+I. INTRODUCTION
+
+This Agreement applies to any individual person or entity ("You", "Your" or "Licensee") that uses or distributes any portion or element of the Stability AI Materials  or Derivative Works thereof for any Research & Non-Commercial or Commercial purpose. Capitalized terms not otherwise defined herein are defined in Section V below.
+
+
+This Agreement is intended to allow research, non-commercial, and limited commercial uses of the Models free of charge. In order to ensure that certain limited commercial uses of the Models continue to be allowed, this Agreement  preserves free access to the Models for people or organizations  generating annual revenue of less than US $1,000,000 (or local currency equivalent).
+
+
+By clicking "I Accept"  or by using or distributing or using any portion or element of the Stability Materials or Derivative Works, You agree that You have read, understood and are bound by the terms of this Agreement. If You are acting on behalf of a company, organization or other entity, then "You" includes you and that entity, and You agree that You: (i) are an authorized representative of such entity with the authority to bind such entity to this Agreement, and (ii) You agree to the terms of this Agreement on that entity's behalf.
+
+II. RESEARCH & NON-COMMERCIAL USE LICENSE
+
+Subject to the terms of this Agreement, Stability AI grants You a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable and royalty-free limited license under Stability AI's intellectual property or other rights owned by Stability AI embodied in the Stability AI Materials to use, reproduce, distribute, and create Derivative Works of, and make modifications to, the Stability AI Materials for any Research or Non-Commercial Purpose. "Research Purpose" means academic or scientific advancement, and in each case, is not primarily intended for commercial advantage or monetary compensation to You or others. "Non-Commercial Purpose" means any purpose other than a Research Purpose that is not primarily intended for commercial advantage or monetary compensation to You or others, such as personal use (i.e., hobbyist) or evaluation and testing.
+
+III. COMMERCIAL USE LICENSE
+
+Subject to the terms of this Agreement (including the remainder of this Section III), Stability AI grants You a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable and royalty-free limited license under Stability AI's intellectual property or other rights owned by Stability AI embodied in the Stability AI Materials to use, reproduce, distribute, and create Derivative Works of, and make modifications to, the Stability AI Materials for any Commercial Purpose. "Commercial Purpose" means any purpose other than a Research Purpose or Non-Commercial Purpose that is primarily intended for commercial advantage or monetary compensation to You or others, including but not limited to, (i) creating, modifying, or distributing Your product or service, including via a hosted service or application programming interface, and (ii) for Your business's or organization's internal operations.
+If You are using or distributing the Stability AI Materials for a Commercial Purpose, You must register with Stability AI at (https://stability.ai/community-license). If at any time You or Your Affiliate(s), either individually or in aggregate, generate more than USD $1,000,000 in annual revenue (or the equivalent thereof in Your local currency), regardless of whether that revenue is generated directly or indirectly from the Stability AI Materials or Derivative Works, any licenses granted to You under this Agreement shall terminate as of such date. You must request a license from Stability AI at (https://stability.ai/enterprise) , which Stability AI may grant to You in its sole discretion. If you receive Stability AI Materials, or any Derivative Works thereof, from a Licensee as part of an integrated end user product, then Section III of this Agreement will not apply to you.
+
+IV. GENERAL TERMS
+
+Your Research, Non-Commercial, and Commercial License(s) under this Agreement are subject to the following terms.
+a.  Distribution & Attribution. If You distribute or make available the Stability AI Materials or a Derivative Work to a third party, or a product or service that uses any portion of them, You shall: (i) provide a copy of this Agreement to that third party, (ii) retain the following attribution notice within a "Notice" text file distributed as a part of such copies: "This Stability AI Model is licensed under the Stability AI Community License, Copyright ©  Stability AI Ltd. All Rights Reserved", and (iii) prominently display "Powered by Stability AI" on a related website, user interface, blogpost, about page, or product documentation.  If You create a Derivative Work, You may add your own attribution notice(s) to the "Notice" text file included with that Derivative Work, provided that You clearly indicate which attributions apply to the Stability AI Materials and state in the "Notice" text file that You changed the Stability AI Materials and how it was modified.
+b.  Use Restrictions. Your use of the Stability AI Materials and Derivative Works, including any output or results of the Stability AI Materials or Derivative Works, must comply with applicable laws and regulations (including Trade Control Laws and equivalent regulations) and adhere to the Documentation and Stability AI's AUP, which is hereby incorporated by reference. Furthermore, You will not use the Stability AI Materials or Derivative Works, or any output or results of the Stability AI Materials or Derivative Works, to create or improve any foundational generative AI model (excluding the Models or Derivative Works).
+c.  Intellectual Property.
+(i) Trademark License.  No trademark licenses are granted under this Agreement, and in connection with the Stability AI Materials or Derivative Works, You may not use any name or mark owned by or associated with Stability AI or any of its Affiliates, except as required under Section IV(a) herein.
+(ii)  Ownership of Derivative Works.  As between You and Stability AI, You are the owner of Derivative Works You create, subject to Stability AI's ownership of the Stability AI Materials and any Derivative Works made by or for Stability AI.
+(iii)  Ownership of Outputs. As between You and Stability AI, You own any outputs generated from the Models or Derivative Works to the extent permitted by applicable law.
+(iv)  Disputes.  If You or Your Affiliate(s) institute litigation or other proceedings against Stability AI (including a cross-claim or counterclaim in a lawsuit) alleging that the Stability AI Materials, Derivative Works or associated outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by You, then any licenses granted to You under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Stability AI from and against any claim by any third party arising out of or related to Your use or distribution of the Stability AI Materials or Derivative Works in violation of this Agreement.
+(v)  Feedback.  From time to time, You may provide Stability AI with verbal and/or written suggestions, comments or other feedback related to Stability AI's existing or prospective technology, products or services (collectively, "Feedback"). You are not obligated to provide Stability AI with Feedback, but to the extent that You do, You hereby grant Stability AI a perpetual, irrevocable, royalty-free, fully-paid, sub-licensable, transferable, non-exclusive, worldwide right and license to exploit the Feedback in any manner without restriction. Your Feedback is provided "AS IS" and You make no warranties whatsoever about any Feedback.
+d.  Disclaimer Of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE STABILITY AI MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OR LAWFULNESS OF USING OR REDISTRIBUTING THE STABILITY AI MATERIALS, DERIVATIVE WORKS OR ANY OUTPUT OR RESULTS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE STABILITY AI MATERIALS, DERIVATIVE WORKS AND ANY OUTPUT AND RESULTS.
+e.  Limitation Of Liability. IN NO EVENT WILL STABILITY AI OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT, INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF STABILITY AI OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+f.  Term And Termination. The term of this Agreement will commence upon Your acceptance of this Agreement or access to the Stability AI Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Stability AI may terminate this Agreement if You are in breach of any term or condition of this Agreement. Upon termination of this Agreement, You shall delete and cease use of any Stability AI Materials or Derivative Works. Section IV(d), (e), and (g) shall survive the termination of this Agreement.
+g.  Governing Law.  This Agreement will be governed by and constructed in accordance with the laws of the United States and the State of California without regard to choice of law principles, and the UN Convention on Contracts for International Sale of Goods does not apply to this Agreement.
+
+V. DEFINITIONS
+
+"Affiliate(s)" means any entity that directly or indirectly controls, is controlled by, or is under common control with the subject entity; for purposes of this definition, "control" means direct or indirect ownership or control of more than 50% of the voting interests of the subject entity.
+"Agreement" means this Stability AI Community License Agreement.
+"AUP" means the Stability AI Acceptable Use Policy available at https://stability.ai/use-policy, as may be updated from time to time.
+"Derivative Work(s)" means (a) any derivative work of the Stability AI Materials as recognized by U.S. copyright laws and (b) any modifications to a Model, and any other model created which is based on or derived from the Model or the Model's output, including"fine tune" and "low-rank adaptation" models derived from a Model or a Model's output, but do not include the output of any Model.
+"Documentation" means any specifications, manuals, documentation, and other written information provided by Stability AI related to the Software or Models.
+"Model(s)" means, collectively, Stability AI's proprietary models and algorithms, including machine-learning models, trained model weights and other elements of the foregoing listed on Stability's Core Models Webpage available at, https://stability.ai/core-models, as may be updated from time to time.
+"Stability AI" or "we" means Stability AI Ltd. and its Affiliates.
+"Software" means Stability AI's proprietary software made available under this Agreement now or in the future.
+"Stability AI Materials" means, collectively, Stability's proprietary Models, Software and Documentation (and any portion or combination thereof) made available under this Agreement.
+"Trade Control Laws" means any applicable U.S. and non-U.S. export control and trade sanctions laws and regulations.

README.md

@@ -1,11 +1,11 @@
 ---
-title: Stable Point Aware 3d
+title: Stable Point-Aware 3D
 emoji: ⚡
 colorFrom: yellow
 colorTo: yellow
 sdk: gradio
-sdk_version: 5.9.1
-app_file: app.py
+sdk_version: 4.43.0
+app_file: gradio_app.py
 pinned: false
 ---
 

__init__.py

@@ -0,0 +1,358 @@
+import base64
+import logging
+import os
+import random
+import sys
+
+import comfy.model_management
+import folder_paths
+import numpy as np
+import torch
+import trimesh
+from PIL import Image
+from trimesh.exchange import gltf
+
+sys.path.append(os.path.dirname(__file__))
+from spar3d.models.mesh import QUAD_REMESH_AVAILABLE, TRIANGLE_REMESH_AVAILABLE
+from spar3d.system import SPAR3D
+from spar3d.utils import foreground_crop
+
+SPAR3D_CATEGORY = "SPAR3D"
+SPAR3D_MODEL_NAME = "stabilityai/spar3d"
+
+
+class SPAR3DLoader:
+    CATEGORY = SPAR3D_CATEGORY
+    FUNCTION = "load"
+    RETURN_NAMES = ("spar3d_model",)
+    RETURN_TYPES = ("SPAR3D_MODEL",)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {}}
+
+    def load(self):
+        device = comfy.model_management.get_torch_device()
+        model = SPAR3D.from_pretrained(
+            SPAR3D_MODEL_NAME,
+            config_name="config.yaml",
+            weight_name="model.safetensors",
+        )
+        model.to(device)
+        model.eval()
+
+        return (model,)
+
+
+class SPAR3DPreview:
+    CATEGORY = SPAR3D_CATEGORY
+    FUNCTION = "preview"
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"mesh": ("MESH",)}}
+
+    def preview(self, mesh):
+        glbs = []
+        for m in mesh:
+            scene = trimesh.Scene(m)
+            glb_data = gltf.export_glb(scene, include_normals=True)
+            glb_base64 = base64.b64encode(glb_data).decode("utf-8")
+            glbs.append(glb_base64)
+        return {"ui": {"glbs": glbs}}
+
+
+class SPAR3DSampler:
+    CATEGORY = SPAR3D_CATEGORY
+    FUNCTION = "predict"
+    RETURN_NAMES = ("mesh", "pointcloud")
+    RETURN_TYPES = ("MESH", "POINTCLOUD")
+
+    @classmethod
+    def INPUT_TYPES(s):
+        remesh_choices = ["none"]
+        if TRIANGLE_REMESH_AVAILABLE:
+            remesh_choices.append("triangle")
+        if QUAD_REMESH_AVAILABLE:
+            remesh_choices.append("quad")
+
+        opt_dict = {
+            "mask": ("MASK",),
+            "pointcloud": ("POINTCLOUD",),
+            "target_type": (["none", "vertex", "face"],),
+            "target_count": (
+                "INT",
+                {"default": 1000, "min": 3, "max": 20000, "step": 1},
+            ),
+            "guidance_scale": (
+                "FLOAT",
+                {"default": 3.0, "min": 1.0, "max": 5.0, "step": 0.05},
+            ),
+            "seed": (
+                "INT",
+                {"default": 42, "min": 0, "max": 2**32 - 1, "step": 1},
+            ),
+        }
+        if TRIANGLE_REMESH_AVAILABLE or QUAD_REMESH_AVAILABLE:
+            opt_dict["remesh"] = (remesh_choices,)
+
+        return {
+            "required": {
+                "model": ("SPAR3D_MODEL",),
+                "image": ("IMAGE",),
+                "foreground_ratio": (
+                    "FLOAT",
+                    {"default": 1.3, "min": 1.0, "max": 2.0, "step": 0.01},
+                ),
+                "texture_resolution": (
+                    "INT",
+                    {"default": 1024, "min": 512, "max": 2048, "step": 256},
+                ),
+            },
+            "optional": opt_dict,
+        }
+
+    def predict(
+        s,
+        model,
+        image,
+        mask,
+        foreground_ratio,
+        texture_resolution,
+        pointcloud=None,
+        remesh="none",
+        target_type="none",
+        target_count=1000,
+        guidance_scale=3.0,
+        seed=42,
+    ):
+        if image.shape[0] != 1:
+            raise ValueError("Only one image can be processed at a time")
+
+        vertex_count = (
+            -1
+            if target_type == "none"
+            else (target_count // 2 if target_type == "face" else target_count)
+        )
+
+        pil_image = Image.fromarray(
+            torch.clamp(torch.round(255.0 * image[0]), 0, 255)
+            .type(torch.uint8)
+            .cpu()
+            .numpy()
+        )
+
+        if mask is not None:
+            print("Using Mask")
+            mask_np = np.clip(255.0 * mask[0].detach().cpu().numpy(), 0, 255).astype(
+                np.uint8
+            )
+            mask_pil = Image.fromarray(mask_np, mode="L")
+            pil_image.putalpha(mask_pil)
+        else:
+            if image.shape[3] != 4:
+                print("No mask or alpha channel detected, Converting to RGBA")
+                pil_image = pil_image.convert("RGBA")
+
+        pil_image = foreground_crop(pil_image, foreground_ratio)
+
+        model.cfg.guidance_scale = guidance_scale
+        random.seed(seed)
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+
+        print(remesh)
+        with torch.no_grad():
+            with torch.autocast(device_type="cuda", dtype=torch.float16):
+                if not TRIANGLE_REMESH_AVAILABLE and remesh == "triangle":
+                    raise ImportError(
+                        "Triangle remeshing requires gpytoolbox to be installed"
+                    )
+                if not QUAD_REMESH_AVAILABLE and remesh == "quad":
+                    raise ImportError("Quad remeshing requires pynim to be installed")
+                mesh, glob_dict = model.run_image(
+                    pil_image,
+                    bake_resolution=texture_resolution,
+                    pointcloud=pointcloud,
+                    remesh=remesh,
+                    vertex_count=vertex_count,
+                )
+
+        if mesh.vertices.shape[0] == 0:
+            raise ValueError("No subject detected in the image")
+
+        return (
+            [mesh],
+            glob_dict["pointcloud"].view(-1).detach().cpu().numpy().tolist(),
+        )
+
+
+class SPAR3DSave:
+    CATEGORY = SPAR3D_CATEGORY
+    FUNCTION = "save"
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mesh": ("MESH",),
+                "filename_prefix": ("STRING", {"default": "SPAR3D"}),
+            }
+        }
+
+    def __init__(self):
+        self.type = "output"
+
+    def save(self, mesh, filename_prefix):
+        output_dir = folder_paths.get_output_directory()
+        glbs = []
+        for idx, m in enumerate(mesh):
+            scene = trimesh.Scene(m)
+            glb_data = gltf.export_glb(scene, include_normals=True)
+            logging.info(f"Generated GLB model with {len(glb_data)} bytes")
+
+            full_output_folder, filename, counter, subfolder, filename_prefix = (
+                folder_paths.get_save_image_path(filename_prefix, output_dir)
+            )
+            filename = filename.replace("%batch_num%", str(idx))
+            out_path = os.path.join(full_output_folder, f"{filename}_{counter:05}_.glb")
+            with open(out_path, "wb") as f:
+                f.write(glb_data)
+            glbs.append(base64.b64encode(glb_data).decode("utf-8"))
+        return {"ui": {"glbs": glbs}}
+
+
+class SPAR3DPointCloudLoader:
+    CATEGORY = SPAR3D_CATEGORY
+    FUNCTION = "load_pointcloud"
+    RETURN_TYPES = ("POINTCLOUD",)
+    RETURN_NAMES = ("pointcloud",)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "file": ("STRING", {"default": None}),
+            }
+        }
+
+    def load_pointcloud(self, file):
+        if file is None or file == "":
+            return (None,)
+        # Load the mesh using trimesh
+        mesh = trimesh.load(file)
+
+        # Extract vertices and colors
+        vertices = mesh.vertices
+
+        # Get vertex colors, defaulting to white if none exist
+        if mesh.visual.vertex_colors is not None:
+            colors = (
+                mesh.visual.vertex_colors[:, :3] / 255.0
+            )  # Convert 0-255 to 0-1 range
+        else:
+            colors = np.ones((len(vertices), 3))
+
+        # Interleave XYZ and RGB values
+        point_cloud = []
+        for vertex, color in zip(vertices, colors):
+            point_cloud.extend(
+                [
+                    float(vertex[0]),
+                    float(vertex[1]),
+                    float(vertex[2]),
+                    float(color[0]),
+                    float(color[1]),
+                    float(color[2]),
+                ]
+            )
+
+        return (point_cloud,)
+
+
+class SPAR3DPointCloudSaver:
+    CATEGORY = SPAR3D_CATEGORY
+    FUNCTION = "save_pointcloud"
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "pointcloud": ("POINTCLOUD",),
+                "filename_prefix": ("STRING", {"default": "SPAR3D"}),
+            }
+        }
+
+    def save_pointcloud(self, pointcloud, filename_prefix):
+        if pointcloud is None:
+            return {"ui": {"text": "No point cloud data to save"}}
+
+        # Reshape the flat list into points with XYZ and RGB
+        points = np.array(pointcloud).reshape(-1, 6)
+
+        # Create vertex array for PLY
+        vertex_array = np.zeros(
+            len(points),
+            dtype=[
+                ("x", "f4"),
+                ("y", "f4"),
+                ("z", "f4"),
+                ("red", "u1"),
+                ("green", "u1"),
+                ("blue", "u1"),
+            ],
+        )
+
+        # Fill vertex array
+        vertex_array["x"] = points[:, 0]
+        vertex_array["y"] = points[:, 1]
+        vertex_array["z"] = points[:, 2]
+        # Convert RGB from 0-1 to 0-255 range
+        vertex_array["red"] = (points[:, 3] * 255).astype(np.uint8)
+        vertex_array["green"] = (points[:, 4] * 255).astype(np.uint8)
+        vertex_array["blue"] = (points[:, 5] * 255).astype(np.uint8)
+
+        # Create PLY object
+        ply_data = trimesh.PointCloud(
+            vertices=points[:, :3], colors=points[:, 3:] * 255
+        )
+
+        # Save to file
+        output_dir = folder_paths.get_output_directory()
+        full_output_folder, filename, counter, subfolder, filename_prefix = (
+            folder_paths.get_save_image_path(filename_prefix, output_dir)
+        )
+        out_path = os.path.join(full_output_folder, f"{filename}_{counter:05}.ply")
+
+        ply_data.export(out_path)
+
+        return {"ui": {"text": f"Saved point cloud to {out_path}"}}
+
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SPAR3DLoader": "SPAR3D Loader",
+    "SPAR3DPreview": "SPAR3D Preview",
+    "SPAR3DSampler": "SPAR3D Sampler",
+    "SPAR3DSave": "SPAR3D Save",
+    "SPAR3DPointCloudLoader": "SPAR3D Point Cloud Loader",
+    "SPAR3DPointCloudSaver": "SPAR3D Point Cloud Saver",
+}
+
+NODE_CLASS_MAPPINGS = {
+    "SPAR3DLoader": SPAR3DLoader,
+    "SPAR3DPreview": SPAR3DPreview,
+    "SPAR3DSampler": SPAR3DSampler,
+    "SPAR3DSave": SPAR3DSave,
+    "SPAR3DPointCloudLoader": SPAR3DPointCloudLoader,
+    "SPAR3DPointCloudSaver": SPAR3DPointCloudSaver,
+}
+
+WEB_DIRECTORY = "./comfyui"
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

demo_files/workflows/spar3d_example.json

@@ -0,0 +1,263 @@
+{
+  "last_node_id": 17,
+  "last_link_id": 18,
+  "nodes": [
+    {
+      "id": 10,
+      "type": "SPAR3DLoader",
+      "pos": [
+        52.92446517944336,
+        394.328369140625
+      ],
+      "size": [
+        210,
+        26
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "spar3d_model",
+          "type": "SPAR3D_MODEL",
+          "links": [
+            10
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "SPAR3DLoader"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 13,
+      "type": "LoadImage",
+      "pos": [
+        -43.437347412109375,
+        482.89678955078125
+      ],
+      "size": [
+        315,
+        314
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            11
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": [
+            16
+          ],
+          "slot_index": 1
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "cat1.png",
+        "image"
+      ]
+    },
+    {
+      "id": 16,
+      "type": "InvertMask",
+      "pos": [
+        377.1180419921875,
+        605.384765625
+      ],
+      "size": [
+        210,
+        26
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "mask",
+          "type": "MASK",
+          "link": 16
+        }
+      ],
+      "outputs": [
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": [
+            17
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "InvertMask"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 17,
+      "type": "SPAR3DSave",
+      "pos": [
+        1133.669921875,
+        439.6551513671875
+      ],
+      "size": [
+        315,
+        58
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "mesh",
+          "type": "MESH",
+          "link": 18
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "SPAR3DSave"
+      },
+      "widgets_values": [
+        "SPAR3D"
+      ]
+    },
+    {
+      "id": 11,
+      "type": "SPAR3DSampler",
+      "pos": [
+        673.0637817382812,
+        441.2229309082031
+      ],
+      "size": [
+        315,
+        286
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "model",
+          "type": "SPAR3D_MODEL",
+          "link": 10
+        },
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 11
+        },
+        {
+          "name": "mask",
+          "type": "MASK",
+          "link": 17,
+          "shape": 7
+        },
+        {
+          "name": "pointcloud",
+          "type": "POINTCLOUD",
+          "link": null,
+          "shape": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "mesh",
+          "type": "MESH",
+          "links": [
+            18
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "pointcloud",
+          "type": "POINTCLOUD",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "SPAR3DSampler"
+      },
+      "widgets_values": [
+        1.3,
+        1024,
+        "none",
+        1000,
+        3,
+        3727502160,
+        "randomize",
+        "none"
+      ]
+    }
+  ],
+  "links": [
+    [
+      10,
+      10,
+      0,
+      11,
+      0,
+      "SPAR3D_MODEL"
+    ],
+    [
+      11,
+      13,
+      0,
+      11,
+      1,
+      "IMAGE"
+    ],
+    [
+      16,
+      13,
+      1,
+      16,
+      0,
+      "MASK"
+    ],
+    [
+      17,
+      16,
+      0,
+      11,
+      2,
+      "MASK"
+    ],
+    [
+      18,
+      11,
+      0,
+      17,
+      0,
+      "MESH"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.953502721998243,
+      "offset": [
+        266.21995970220667,
+        116.75398112171928
+      ]
+    }
+  },
+  "version": 0.4
+}

gradio_app.py

@@ -0,0 +1,792 @@
+import os
+
+os.system("pip install ./texture_baker/ ./uv_unwrapper/")
+
+import random
+import tempfile
+import time
+from contextlib import nullcontext
+from functools import lru_cache
+from typing import Any
+
+import gradio as gr
+import numpy as np
+import torch
+import trimesh
+from gradio_litmodel3d import LitModel3D
+from gradio_pointcloudeditor import PointCloudEditor
+from PIL import Image
+from transparent_background import Remover
+
+import spar3d.utils as spar3d_utils
+from spar3d.models.mesh import QUAD_REMESH_AVAILABLE, TRIANGLE_REMESH_AVAILABLE
+from spar3d.system import SPAR3D
+
+os.environ["GRADIO_TEMP_DIR"] = os.path.join(os.environ.get("TMPDIR", "/tmp"), "gradio")
+
+bg_remover = Remover()  # default setting
+
+COND_WIDTH = 512
+COND_HEIGHT = 512
+COND_DISTANCE = 2.2
+COND_FOVY = 0.591627
+BACKGROUND_COLOR = [0.5, 0.5, 0.5]
+
+# Cached. Doesn't change
+c2w_cond = spar3d_utils.default_cond_c2w(COND_DISTANCE)
+intrinsic, intrinsic_normed_cond = spar3d_utils.create_intrinsic_from_fov_rad(
+    COND_FOVY, COND_HEIGHT, COND_WIDTH
+)
+
+generated_files = []
+
+# Delete previous gradio temp dir folder
+if os.path.exists(os.environ["GRADIO_TEMP_DIR"]):
+    print(f"Deleting {os.environ['GRADIO_TEMP_DIR']}")
+    import shutil
+
+    shutil.rmtree(os.environ["GRADIO_TEMP_DIR"])
+
+device = spar3d_utils.get_device()
+
+model = SPAR3D.from_pretrained(
+    "stabilityai/stable-point-aware-3d",
+    config_name="config.yaml",
+    weight_name="model.safetensors",
+)
+model.eval()
+model = model.to(device)
+
+example_files = [
+    os.path.join("demo_files/examples", f) for f in os.listdir("demo_files/examples")
+]
+
+
+def forward_model(
+    batch,
+    system,
+    guidance_scale=3.0,
+    seed=0,
+    device="cuda",
+    remesh_option="none",
+    vertex_count=-1,
+    texture_resolution=1024,
+):
+    batch_size = batch["rgb_cond"].shape[0]
+
+    # prepare the condition for point cloud generation
+    # set seed
+    random.seed(seed)
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    cond_tokens = system.forward_pdiff_cond(batch)
+
+    if "pc_cond" not in batch:
+        sample_iter = system.sampler.sample_batch_progressive(
+            batch_size,
+            cond_tokens,
+            guidance_scale=guidance_scale,
+            device=device,
+        )
+        for x in sample_iter:
+            samples = x["xstart"]
+        batch["pc_cond"] = samples.permute(0, 2, 1).float()
+        batch["pc_cond"] = spar3d_utils.normalize_pc_bbox(batch["pc_cond"])
+
+    # subsample to the 512 points
+    batch["pc_cond"] = batch["pc_cond"][
+        :, torch.randperm(batch["pc_cond"].shape[1])[:512]
+    ]
+
+    # get the point cloud
+    xyz = batch["pc_cond"][0, :, :3].cpu().numpy()
+    color_rgb = (batch["pc_cond"][0, :, 3:6] * 255).cpu().numpy().astype(np.uint8)
+    pc_rgb_trimesh = trimesh.PointCloud(vertices=xyz, colors=color_rgb)
+
+    # forward for the final mesh
+    trimesh_mesh, _glob_dict = model.generate_mesh(
+        batch, texture_resolution, remesh=remesh_option, vertex_count=vertex_count
+    )
+    trimesh_mesh = trimesh_mesh[0]
+
+    return trimesh_mesh, pc_rgb_trimesh
+
+
+def run_model(
+    input_image,
+    guidance_scale,
+    random_seed,
+    pc_cond,
+    remesh_option,
+    vertex_count,
+    texture_resolution,
+):
+    start = time.time()
+    with torch.no_grad():
+        with (
+            torch.autocast(device_type=device, dtype=torch.float16)
+            if "cuda" in device
+            else nullcontext()
+        ):
+            model_batch = create_batch(input_image)
+            model_batch = {k: v.to(device) for k, v in model_batch.items()}
+
+            if pc_cond is not None:
+                # Check if pc_cond is a list
+                if isinstance(pc_cond, list):
+                    cond_tensor = torch.tensor(pc_cond).float().cuda().view(-1, 6)
+                    xyz = cond_tensor[:, :3]
+                    color_rgb = cond_tensor[:, 3:]
+                elif isinstance(pc_cond, dict):
+                    xyz = torch.tensor(pc_cond["positions"]).float().cuda()
+                    color_rgb = torch.tensor(pc_cond["colors"]).float().cuda()
+                else:
+                    xyz = torch.tensor(pc_cond.vertices).float().cuda()
+                    color_rgb = (
+                        torch.tensor(pc_cond.colors[:, :3]).float().cuda() / 255.0
+                    )
+                model_batch["pc_cond"] = torch.cat([xyz, color_rgb], dim=-1).unsqueeze(
+                    0
+                )
+                # sub-sample the point cloud to the target number of points
+                if model_batch["pc_cond"].shape[1] > 512:
+                    idx = torch.randperm(model_batch["pc_cond"].shape[1])[:512]
+                    model_batch["pc_cond"] = model_batch["pc_cond"][:, idx]
+                elif model_batch["pc_cond"].shape[1] < 512:
+                    num_points = model_batch["pc_cond"].shape[1]
+                    gr.Warning(
+                        f"The uploaded point cloud should have at least 512 points. This point cloud only has {num_points}. Results may be worse."
+                    )
+                    pad = 512 - num_points
+                    sampled_idx = torch.randint(
+                        0, model_batch["pc_cond"].shape[1], (pad,)
+                    )
+                    model_batch["pc_cond"] = torch.cat(
+                        [
+                            model_batch["pc_cond"],
+                            model_batch["pc_cond"][:, sampled_idx],
+                        ],
+                        dim=1,
+                    )
+
+            trimesh_mesh, trimesh_pc = forward_model(
+                model_batch,
+                model,
+                guidance_scale=guidance_scale,
+                seed=random_seed,
+                device="cuda",
+                remesh_option=remesh_option.lower(),
+                vertex_count=vertex_count,
+                texture_resolution=texture_resolution,
+            )
+
+    # Create new tmp file
+    temp_dir = tempfile.mkdtemp()
+    tmp_file = os.path.join(temp_dir, "mesh.glb")
+
+    trimesh_mesh.export(tmp_file, file_type="glb", include_normals=True)
+    generated_files.append(tmp_file)
+
+    tmp_file_pc = os.path.join(temp_dir, "points.ply")
+    trimesh_pc.export(tmp_file_pc)
+    generated_files.append(tmp_file_pc)
+
+    print("Generation took:", time.time() - start, "s")
+
+    return tmp_file, tmp_file_pc, trimesh_pc
+
+
+def create_batch(input_image: Image) -> dict[str, Any]:
+    img_cond = (
+        torch.from_numpy(
+            np.asarray(input_image.resize((COND_WIDTH, COND_HEIGHT))).astype(np.float32)
+            / 255.0
+        )
+        .float()
+        .clip(0, 1)
+    )
+    mask_cond = img_cond[:, :, -1:]
+    rgb_cond = torch.lerp(
+        torch.tensor(BACKGROUND_COLOR)[None, None, :], img_cond[:, :, :3], mask_cond
+    )
+
+    batch_elem = {
+        "rgb_cond": rgb_cond,
+        "mask_cond": mask_cond,
+        "c2w_cond": c2w_cond.unsqueeze(0),
+        "intrinsic_cond": intrinsic.unsqueeze(0),
+        "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),
+    }
+    # Add batch dim
+    batched = {k: v.unsqueeze(0) for k, v in batch_elem.items()}
+    return batched
+
+
+@lru_cache
+def checkerboard(squares: int, size: int, min_value: float = 0.5):
+    base = np.zeros((squares, squares)) + min_value
+    base[1::2, ::2] = 1
+    base[::2, 1::2] = 1
+
+    repeat_mult = size // squares
+    return (
+        base.repeat(repeat_mult, axis=0)
+        .repeat(repeat_mult, axis=1)[:, :, None]
+        .repeat(3, axis=-1)
+    )
+
+
+def remove_background(input_image: Image) -> Image:
+    return bg_remover.process(input_image.convert("RGB"))
+
+
+def show_mask_img(input_image: Image) -> Image:
+    img_numpy = np.array(input_image)
+    alpha = img_numpy[:, :, 3] / 255.0
+    chkb = checkerboard(32, 512) * 255
+    new_img = img_numpy[..., :3] * alpha[:, :, None] + chkb * (1 - alpha[:, :, None])
+    return Image.fromarray(new_img.astype(np.uint8), mode="RGB")
+
+
+def process_model_run(
+    background_state,
+    guidance_scale,
+    random_seed,
+    pc_cond,
+    remesh_option,
+    vertex_count_type,
+    vertex_count,
+    texture_resolution,
+):
+    # Adjust vertex count based on selection
+    final_vertex_count = (
+        -1
+        if vertex_count_type == "Keep Vertex Count"
+        else (
+            vertex_count // 2
+            if vertex_count_type == "Target Face Count"
+            else vertex_count
+        )
+    )
+    print(
+        f"Final vertex count: {final_vertex_count} with type {vertex_count_type} and vertex count {vertex_count}"
+    )
+
+    glb_file, pc_file, pc_plot = run_model(
+        background_state,
+        guidance_scale,
+        random_seed,
+        pc_cond,
+        remesh_option,
+        final_vertex_count,
+        texture_resolution,
+    )
+    # Create a single float list of x y z r g b
+    point_list = []
+    for i in range(pc_plot.vertices.shape[0]):
+        point_list.extend(
+            [
+                pc_plot.vertices[i, 0],
+                pc_plot.vertices[i, 1],
+                pc_plot.vertices[i, 2],
+                pc_plot.colors[i, 0] / 255.0,
+                pc_plot.colors[i, 1] / 255.0,
+                pc_plot.colors[i, 2] / 255.0,
+            ]
+        )
+
+    return glb_file, pc_file, point_list
+
+
+def regenerate_run(
+    background_state,
+    guidance_scale,
+    random_seed,
+    pc_cond,
+    remesh_option,
+    vertex_count_type,
+    vertex_count,
+    texture_resolution,
+):
+    glb_file, pc_file, point_list = process_model_run(
+        background_state,
+        guidance_scale,
+        random_seed,
+        pc_cond,
+        remesh_option,
+        vertex_count_type,
+        vertex_count,
+        texture_resolution,
+    )
+    return (
+        gr.update(),  # run_btn
+        gr.update(),  # img_proc_state
+        gr.update(),  # background_remove_state
+        gr.update(),  # preview_removal
+        gr.update(value=glb_file, visible=True),  # output_3d
+        gr.update(visible=True),  # hdr_row
+        gr.update(visible=True),  # point_cloud_row
+        gr.update(value=point_list),  # point_cloud_editor
+        gr.update(value=pc_file),  # pc_download
+        gr.update(visible=False),  # regenerate_btn
+    )
+
+
+def run_button(
+    run_btn,
+    input_image,
+    background_state,
+    foreground_ratio,
+    no_crop,
+    guidance_scale,
+    random_seed,
+    pc_upload,
+    pc_cond_file,
+    remesh_option,
+    vertex_count_type,
+    vertex_count,
+    texture_resolution,
+):
+    if run_btn == "Run":
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+
+        if pc_upload:
+            # make sure the pc_cond_file has been uploaded
+            try:
+                pc_cond = trimesh.load(pc_cond_file.name)
+            except Exception:
+                raise gr.Error(
+                    "Please upload a valid point cloud ply file as condition."
+                )
+        else:
+            pc_cond = None
+
+        glb_file, pc_file, pc_list = process_model_run(
+            background_state,
+            guidance_scale,
+            random_seed,
+            pc_cond,
+            remesh_option,
+            vertex_count_type,
+            vertex_count,
+            texture_resolution,
+        )
+
+        if torch.cuda.is_available():
+            print("Peak Memory:", torch.cuda.max_memory_allocated() / 1024 / 1024, "MB")
+        elif torch.backends.mps.is_available():
+            print(
+                "Peak Memory:", torch.mps.driver_allocated_memory() / 1024 / 1024, "MB"
+            )
+
+        return (
+            gr.update(),  # run_btn
+            gr.update(),  # img_proc_state
+            gr.update(),  # background_remove_state
+            gr.update(),  # preview_removal
+            gr.update(value=glb_file, visible=True),  # output_3d
+            gr.update(visible=True),  # hdr_row
+            gr.update(visible=True),  # point_cloud_row
+            gr.update(value=pc_list),  # point_cloud_editor
+            gr.update(value=pc_file),  # pc_download
+            gr.update(visible=False),  # regenerate_btn
+        )
+
+    elif run_btn == "Remove Background":
+        rem_removed = remove_background(input_image)
+
+        fr_res = spar3d_utils.foreground_crop(
+            rem_removed,
+            crop_ratio=foreground_ratio,
+            newsize=(COND_WIDTH, COND_HEIGHT),
+            no_crop=no_crop,
+        )
+
+        return (
+            gr.update(value="Run", visible=True),  # run_btn
+            rem_removed,  # img_proc_state,
+            fr_res,  # background_remove_state
+            gr.update(value=show_mask_img(fr_res), visible=True),  # preview_removal
+            gr.update(value=None, visible=False),  # output_3d
+            gr.update(visible=False),  # hdr_row
+            gr.update(visible=False),  # point_cloud_row
+            gr.update(value=None),  # point_cloud_editor
+            gr.update(value=None),  # pc_download
+            gr.update(visible=False),  # regenerate_btn
+        )
+
+
+def requires_bg_remove(image, fr, no_crop):
+    if image is None:
+        return (
+            gr.update(visible=False, value="Run"),  # run_Btn
+            None,  # img_proc_state
+            None,  # background_remove_state
+            gr.update(value=None, visible=False),  # preview_removal
+            gr.update(value=None, visible=False),  # output_3d
+            gr.update(visible=False),  # hdr_row
+            gr.update(visible=False),  # point_cloud_row
+            gr.update(value=None),  # point_cloud_editor
+            gr.update(value=None),  # pc_download
+            gr.update(visible=False),  # regenerate_btn
+        )
+    alpha_channel = np.array(image.getchannel("A"))
+    min_alpha = alpha_channel.min()
+
+    if min_alpha == 0:
+        print("Already has alpha")
+        fr_res = spar3d_utils.foreground_crop(
+            image, fr, newsize=(COND_WIDTH, COND_HEIGHT), no_crop=no_crop
+        )
+        return (
+            gr.update(value="Run", visible=True),  # run_Btn
+            image,  # img_proc_state
+            fr_res,  # background_remove_state
+            gr.update(value=show_mask_img(fr_res), visible=True),  # preview_removal
+            gr.update(value=None, visible=False),  # output_3d
+            gr.update(visible=False),  # hdr_row
+            gr.update(visible=False),  # point_cloud_row
+            gr.update(value=None),  # point_cloud_editor
+            gr.update(value=None),  # pc_download
+            gr.update(visible=False),  # regenerate_btn
+        )
+    return (
+        gr.update(value="Remove Background", visible=True),  # run_Btn
+        None,  # img_proc_state
+        None,  # background_remove_state
+        gr.update(value=None, visible=False),  # preview_removal
+        gr.update(value=None, visible=False),  # output_3d
+        gr.update(visible=False),  # hdr_row
+        gr.update(visible=False),  # point_cloud_row
+        gr.update(value=None),  # point_cloud_editor
+        gr.update(value=None),  # pc_download
+        gr.update(visible=False),  # regenerate_btn
+    )
+
+
+def update_foreground_ratio(img_proc, fr, no_crop):
+    foreground_res = spar3d_utils.foreground_crop(
+        img_proc, fr, newsize=(COND_WIDTH, COND_HEIGHT), no_crop=no_crop
+    )
+    return (
+        foreground_res,
+        gr.update(value=show_mask_img(foreground_res)),
+    )
+
+
+def update_resolution_controls(remesh_choice, vertex_count_type):
+    show_controls = remesh_choice.lower() != "none"
+    show_vertex_count = vertex_count_type != "Keep Vertex Count"
+    return (
+        gr.update(visible=show_controls),  # vertex_count_type
+        gr.update(visible=show_controls and show_vertex_count),  # vertex_count_slider
+    )
+
+
+with gr.Blocks() as demo:
+    img_proc_state = gr.State()
+    background_remove_state = gr.State()
+    gr.Markdown(
+        """
+    # SPAR3D: Stable Point-Aware Reconstruction of 3D Objects from Single Images
+
+    SPAR3D is a state-of-the-art method for 3D mesh reconstruction from a single image. This demo allows you to upload an image and generate a 3D mesh model from it. A feature of SPAR3D is it generates point clouds as intermediate representation before producing the mesh. You can edit the point cloud to adjust the final mesh. We provide a simple point cloud editor in this demo, where you can drag, recolor and rescale the point clouds. If you have more advanced editing needs (e.g. box selection, duplication, local streching, etc.), you can download the point cloud and edit it in softwares such as MeshLab or Blender. The edited point cloud can then be uploaded to this demo to generate a new 3D model by checking the "Point cloud upload" box.
+
+    **Tips**
+
+    1. If the image does not have a valid alpha channel, it will go through the background removal step. Our built-in background removal can be inaccurate sometimes, which will result in poor mesh quality. In such cases, you can use external background removal tools to obtain a RGBA image before uploading here.
+    2. You can adjust the foreground ratio to control the size of the foreground object. This may have major impact on the final mesh.
+    3. Guidance scale controls the strength of the image condition in the point cloud generation process. A higher value may result in higher mesh fidelity, but the variability by changing the random seed will be lower. Note that the guidance scale and the seed are not effective when the point cloud is manually uploaded.
+    4. Our online editor supports multi-selection by holding down the shift key. This allows you to recolor multiple points at once.
+    5. The editing should mainly alter the unseen parts of the object. Visible parts can be edited, but the edits should be consistent with the image. Editing the visible parts in a way that contradicts the image may result in poor mesh quality.
+    6. You can upload your own HDR environment map to light the 3D model.
+    """
+    )
+    with gr.Row(variant="panel"):
+        with gr.Column():
+            with gr.Row():
+                input_img = gr.Image(
+                    type="pil", label="Input Image", sources="upload", image_mode="RGBA"
+                )
+                preview_removal = gr.Image(
+                    label="Preview Background Removal",
+                    type="pil",
+                    image_mode="RGB",
+                    interactive=False,
+                    visible=False,
+                )
+
+            gr.Markdown("### Input Controls")
+            with gr.Group():
+                with gr.Row():
+                    no_crop = gr.Checkbox(label="No cropping", value=False)
+                    pc_upload = gr.Checkbox(label="Point cloud upload", value=False)
+
+                pc_cond_file = gr.File(
+                    label="Point Cloud Upload",
+                    file_types=[".ply"],
+                    file_count="single",
+                    visible=False,
+                )
+
+                foreground_ratio = gr.Slider(
+                    label="Padding Ratio",
+                    minimum=1.0,
+                    maximum=2.0,
+                    value=1.3,
+                    step=0.05,
+                )
+
+            pc_upload.change(
+                lambda x: gr.update(visible=x),
+                inputs=pc_upload,
+                outputs=[pc_cond_file],
+            )
+
+            no_crop.change(
+                update_foreground_ratio,
+                inputs=[img_proc_state, foreground_ratio, no_crop],
+                outputs=[background_remove_state, preview_removal],
+            )
+
+            foreground_ratio.change(
+                update_foreground_ratio,
+                inputs=[img_proc_state, foreground_ratio, no_crop],
+                outputs=[background_remove_state, preview_removal],
+            )
+
+            gr.Markdown("### Point Diffusion Controls")
+            with gr.Group():
+                guidance_scale = gr.Slider(
+                    label="Guidance Scale",
+                    minimum=1.0,
+                    maximum=10.0,
+                    value=3.0,
+                    step=1.0,
+                )
+
+                random_seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=10000,
+                    value=0,
+                    step=1,
+                )
+
+            no_remesh = not TRIANGLE_REMESH_AVAILABLE and not QUAD_REMESH_AVAILABLE
+            gr.Markdown(
+                "### Texture Controls"
+                if no_remesh
+                else "### Meshing and Texture Controls"
+            )
+            with gr.Group():
+                remesh_choices = ["None"]
+                if TRIANGLE_REMESH_AVAILABLE:
+                    remesh_choices.append("Triangle")
+                if QUAD_REMESH_AVAILABLE:
+                    remesh_choices.append("Quad")
+
+                remesh_option = gr.Radio(
+                    choices=remesh_choices,
+                    label="Remeshing",
+                    value="None",
+                    visible=not no_remesh,
+                )
+
+                vertex_count_type = gr.Radio(
+                    choices=[
+                        "Keep Vertex Count",
+                        "Target Vertex Count",
+                        "Target Face Count",
+                    ],
+                    label="Mesh Resolution Control",
+                    value="Keep Vertex Count",
+                    visible=False,
+                )
+
+                vertex_count_slider = gr.Slider(
+                    label="Target Count",
+                    minimum=0,
+                    maximum=20000,
+                    value=2000,
+                    visible=False,
+                )
+
+                texture_size = gr.Slider(
+                    label="Texture Size",
+                    minimum=512,
+                    maximum=2048,
+                    value=1024,
+                    step=256,
+                    visible=True,
+                )
+
+            remesh_option.change(
+                update_resolution_controls,
+                inputs=[remesh_option, vertex_count_type],
+                outputs=[vertex_count_type, vertex_count_slider],
+            )
+
+            vertex_count_type.change(
+                update_resolution_controls,
+                inputs=[remesh_option, vertex_count_type],
+                outputs=[vertex_count_type, vertex_count_slider],
+            )
+
+            run_btn = gr.Button("Run", variant="primary", visible=False)
+
+        with gr.Column():
+            with gr.Group(visible=False) as point_cloud_row:
+                point_size_slider = gr.Slider(
+                    label="Point Size",
+                    minimum=0.01,
+                    maximum=1.0,
+                    value=0.2,
+                    step=0.01,
+                )
+                point_cloud_editor = PointCloudEditor(
+                    up_axis="Z",
+                    forward_axis="X",
+                    lock_scale_z=True,
+                    lock_scale_y=True,
+                    visible=True,
+                )
+
+                pc_download = gr.File(
+                    label="Point Cloud Download",
+                    file_types=[".ply"],
+                    file_count="single",
+                )
+            point_size_slider.change(
+                fn=lambda x: gr.update(point_size=x),
+                inputs=point_size_slider,
+                outputs=point_cloud_editor,
+            )
+
+            regenerate_btn = gr.Button(
+                "Re-run with point cloud", variant="primary", visible=False
+            )
+
+            output_3d = LitModel3D(
+                label="3D Model",
+                visible=False,
+                clear_color=[0.0, 0.0, 0.0, 0.0],
+                tonemapping="aces",
+                contrast=1.0,
+                scale=1.0,
+            )
+            with gr.Column(visible=False, scale=1.0) as hdr_row:
+                gr.Markdown(
+                    """## HDR Environment Map
+
+                Select an HDR environment map to light the 3D model. You can also upload your own HDR environment maps.
+                """
+                )
+
+                with gr.Row():
+                    hdr_illumination_file = gr.File(
+                        label="HDR Env Map",
+                        file_types=[".hdr"],
+                        file_count="single",
+                    )
+                    example_hdris = [
+                        os.path.join("demo_files/hdri", f)
+                        for f in os.listdir("demo_files/hdri")
+                    ]
+                    hdr_illumination_example = gr.Examples(
+                        examples=example_hdris,
+                        inputs=hdr_illumination_file,
+                    )
+
+                    hdr_illumination_file.change(
+                        lambda x: gr.update(env_map=x.name if x is not None else None),
+                        inputs=hdr_illumination_file,
+                        outputs=[output_3d],
+                    )
+
+    examples = gr.Examples(
+        examples=example_files, inputs=input_img, examples_per_page=11
+    )
+
+    input_img.change(
+        requires_bg_remove,
+        inputs=[input_img, foreground_ratio, no_crop],
+        outputs=[
+            run_btn,
+            img_proc_state,
+            background_remove_state,
+            preview_removal,
+            output_3d,
+            hdr_row,
+            point_cloud_row,
+            point_cloud_editor,
+            pc_download,
+            regenerate_btn,
+        ],
+    )
+
+    point_cloud_editor.edit(
+        fn=lambda _x: gr.update(visible=True),
+        inputs=point_cloud_editor,
+        outputs=regenerate_btn,
+    )
+
+    regenerate_btn.click(
+        regenerate_run,
+        inputs=[
+            background_remove_state,
+            guidance_scale,
+            random_seed,
+            point_cloud_editor,
+            remesh_option,
+            vertex_count_type,
+            vertex_count_slider,
+            texture_size,
+        ],
+        outputs=[
+            run_btn,
+            img_proc_state,
+            background_remove_state,
+            preview_removal,
+            output_3d,
+            hdr_row,
+            point_cloud_row,
+            point_cloud_editor,
+            pc_download,
+            regenerate_btn,
+        ],
+    )
+
+    run_btn.click(
+        run_button,
+        inputs=[
+            run_btn,
+            input_img,
+            background_remove_state,
+            foreground_ratio,
+            no_crop,
+            guidance_scale,
+            random_seed,
+            pc_upload,
+            pc_cond_file,
+            remesh_option,
+            vertex_count_type,
+            vertex_count_slider,
+            texture_size,
+        ],
+        outputs=[
+            run_btn,
+            img_proc_state,
+            background_remove_state,
+            preview_removal,
+            output_3d,
+            hdr_row,
+            point_cloud_row,
+            point_cloud_editor,
+            pc_download,
+            regenerate_btn,
+        ],
+    )
+
+demo.queue().launch()

load/tets/160_tets.npz

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

requirements.txt

@@ -0,0 +1,17 @@
+einops==0.7.0
+jaxtyping==0.2.31
+omegaconf==2.3.0
+transformers==4.42.3
+loralib==0.1.2
+git+https://github.com/openai/CLIP.git
+git+https://github.com/SunzeY/AlphaCLIP.git
+trimesh==4.4.1
+numpy==1.26.4
+huggingface-hub==0.23.4
+transparent-background==1.3.3
+gradio==4.43.0
+gradio-litmodel3d==0.0.1
+gradio-pointcloudeditor==0.0.9
+gpytoolbox==0.2.0
+# ./texture_baker/
+# ./uv_unwrapper/

ruff.toml

@@ -0,0 +1,3 @@
+[lint]
+ignore = ["F722", "F821"]
+extend-select = ["I"]

run.py

@@ -0,0 +1,180 @@
+import argparse
+import os
+from contextlib import nullcontext
+
+import torch
+from PIL import Image
+from tqdm import tqdm
+from transparent_background import Remover
+
+from spar3d.models.mesh import QUAD_REMESH_AVAILABLE, TRIANGLE_REMESH_AVAILABLE
+from spar3d.system import SPAR3D
+from spar3d.utils import foreground_crop, get_device, remove_background
+
+
+def check_positive(value):
+    ivalue = int(value)
+    if ivalue <= 0:
+        raise argparse.ArgumentTypeError("%s is an invalid positive int value" % value)
+    return ivalue
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "image", type=str, nargs="+", help="Path to input image(s) or folder."
+    )
+    parser.add_argument(
+        "--device",
+        default=get_device(),
+        type=str,
+        help=f"Device to use. If no CUDA/MPS-compatible device is found, the baking will fail. Default: '{get_device()}'",
+    )
+    parser.add_argument(
+        "--pretrained-model",
+        default="stabilityai/spar3d",
+        type=str,
+        help="Path to the pretrained model. Could be either a huggingface model id is or a local path. Default: 'stabilityai/spar3d'",
+    )
+    parser.add_argument(
+        "--foreground-ratio",
+        default=1.3,
+        type=float,
+        help="Ratio of the foreground size to the image size. Only used when --no-remove-bg is not specified. Default: 0.85",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="output/",
+        type=str,
+        help="Output directory to save the results. Default: 'output/'",
+    )
+    parser.add_argument(
+        "--texture-resolution",
+        default=1024,
+        type=int,
+        help="Texture atlas resolution. Default: 1024",
+    )
+
+    remesh_choices = ["none"]
+    if TRIANGLE_REMESH_AVAILABLE:
+        remesh_choices.append("triangle")
+    if QUAD_REMESH_AVAILABLE:
+        remesh_choices.append("quad")
+    parser.add_argument(
+        "--remesh_option",
+        choices=remesh_choices,
+        default="none",
+        help="Remeshing option",
+    )
+    if TRIANGLE_REMESH_AVAILABLE or QUAD_REMESH_AVAILABLE:
+        parser.add_argument(
+            "--reduction_count_type",
+            choices=["keep", "vertex", "faces"],
+            default="keep",
+            help="Vertex count type",
+        )
+        parser.add_argument(
+            "--target_count",
+            type=check_positive,
+            help="Selected target count.",
+            default=2000,
+        )
+    parser.add_argument(
+        "--batch_size", default=1, type=int, help="Batch size for inference"
+    )
+    args = parser.parse_args()
+
+    # Ensure args.device contains cuda
+    devices = ["cuda", "mps", "cpu"]
+    if not any(args.device in device for device in devices):
+        raise ValueError("Invalid device. Use cuda, mps or cpu")
+
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    device = args.device
+    if not (torch.cuda.is_available() or torch.backends.mps.is_available()):
+        device = "cpu"
+
+    print("Device used: ", device)
+
+    model = SPAR3D.from_pretrained(
+        args.pretrained_model,
+        config_name="config.yaml",
+        weight_name="model.safetensors",
+    )
+    model.to(device)
+    model.eval()
+
+    bg_remover = Remover(device=device)
+    images = []
+    idx = 0
+    for image_path in args.image:
+
+        def handle_image(image_path, idx):
+            image = remove_background(
+                Image.open(image_path).convert("RGBA"), bg_remover
+            )
+            image = foreground_crop(image, args.foreground_ratio)
+            os.makedirs(os.path.join(output_dir, str(idx)), exist_ok=True)
+            image.save(os.path.join(output_dir, str(idx), "input.png"))
+            images.append(image)
+
+        if os.path.isdir(image_path):
+            image_paths = [
+                os.path.join(image_path, f)
+                for f in os.listdir(image_path)
+                if f.endswith((".png", ".jpg", ".jpeg"))
+            ]
+            for image_path in image_paths:
+                handle_image(image_path, idx)
+                idx += 1
+        else:
+            handle_image(image_path, idx)
+            idx += 1
+
+    vertex_count = (
+        -1
+        if args.reduction_count_type == "keep"
+        else (
+            args.target_count
+            if args.reduction_count_type == "vertex"
+            else args.target_count // 2
+        )
+    )
+
+    for i in tqdm(range(0, len(images), args.batch_size)):
+        image = images[i : i + args.batch_size]
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+        with torch.no_grad():
+            with (
+                torch.autocast(device_type=device, dtype=torch.float16)
+                if "cuda" in device
+                else nullcontext()
+            ):
+                mesh, glob_dict = model.run_image(
+                    image,
+                    bake_resolution=args.texture_resolution,
+                    remesh=args.remesh_option,
+                    vertex_count=args.target_vertex_count,
+                    return_points=True,
+                )
+        if torch.cuda.is_available():
+            print("Peak Memory:", torch.cuda.max_memory_allocated() / 1024 / 1024, "MB")
+        elif torch.backends.mps.is_available():
+            print(
+                "Peak Memory:", torch.mps.driver_allocated_memory() / 1024 / 1024, "MB"
+            )
+
+        if len(image) == 1:
+            out_mesh_path = os.path.join(output_dir, str(i), "mesh.glb")
+            mesh.export(out_mesh_path, include_normals=True)
+            out_points_path = os.path.join(output_dir, str(i), "points.ply")
+            glob_dict["point_clouds"][0].export(out_points_path)
+        else:
+            for j in range(len(mesh)):
+                out_mesh_path = os.path.join(output_dir, str(i + j), "mesh.glb")
+                mesh[j].export(out_mesh_path, include_normals=True)
+                out_points_path = os.path.join(output_dir, str(i + j), "points.ply")
+                glob_dict["point_clouds"][j].export(out_points_path)

spar3d/models/camera.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass, field
+from typing import List
+
+import torch
+import torch.nn as nn
+
+from spar3d.models.utils import BaseModule
+
+
+class LinearCameraEmbedder(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int = 25
+        out_channels: int = 768
+        conditions: List[str] = field(default_factory=list)
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.linear = nn.Linear(self.cfg.in_channels, self.cfg.out_channels)
+
+    def forward(self, **kwargs):
+        cond_tensors = []
+        for cond_name in self.cfg.conditions:
+            assert cond_name in kwargs
+            cond = kwargs[cond_name]
+            # cond in shape (B, Nv, ...)
+            cond_tensors.append(cond.view(*cond.shape[:2], -1))
+        cond_tensor = torch.cat(cond_tensors, dim=-1)
+        assert cond_tensor.shape[-1] == self.cfg.in_channels
+        embedding = self.linear(cond_tensor)
+        return embedding

spar3d/models/diffusion/gaussian_diffusion.py

@@ -0,0 +1,524 @@
+# --------------------------------------------------------
+# Adapted from: https://github.com/openai/point-e
+# Licensed under the MIT License
+# Copyright (c) 2022 OpenAI
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# --------------------------------------------------------
+
+import math
+from typing import Any, Dict, Iterable, Optional, Sequence, Union
+
+import numpy as np
+import torch as th
+
+
+def sigmoid_schedule(t, start=-3, end=3, tau=0.6, clip_min=1e-9):
+    def sigmoid(x):
+        return 1 / (1 + np.exp(-x))
+
+    v_start = sigmoid(start / tau)
+    v_end = sigmoid(end / tau)
+    output = sigmoid((t * (end - start) + start) / tau)
+    output = (v_end - output) / (v_end - v_start)
+    return np.clip(output, clip_min, 1.0)
+
+
+def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
+    """
+    This is the deprecated API for creating beta schedules.
+
+    See get_named_beta_schedule() for the new library of schedules.
+    """
+    if beta_schedule == "linear":
+        betas = np.linspace(
+            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
+        )
+    else:
+        raise NotImplementedError(beta_schedule)
+    assert betas.shape == (num_diffusion_timesteps,)
+    return betas
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps, exp_p=12):
+    """
+    Get a pre-defined beta schedule for the given name.
+
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        return get_beta_schedule(
+            "linear",
+            beta_start=scale * 0.0001,
+            beta_end=scale * 0.02,
+            num_diffusion_timesteps=num_diffusion_timesteps,
+        )
+    elif schedule_name == "cosine":
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    elif schedule_name == "sigmoid":
+        # Sigmoid schedule passed through betas_for_alpha_bar
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps, lambda t: sigmoid_schedule(t)
+        )
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim") :])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(
+                f"cannot create exactly {num_timesteps} steps with an integer stride"
+            )
+        elif section_counts.startswith("exact"):
+            res = set(int(x) for x in section_counts[len("exact") :].split(","))
+            for x in res:
+                if x < 0 or x >= num_timesteps:
+                    raise ValueError(f"timestep out of bounds: {x}")
+            return res
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(
+                f"cannot divide section of {size} steps into {section_count}"
+            )
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """Extract values from a 1-D numpy array for a batch of indices."""
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res + th.zeros(broadcast_shape, device=timesteps.device)
+
+
+class GaussianDiffusion:
+    """
+    Utilities for sampling from Gaussian diffusion models.
+    """
+
+    def __init__(
+        self,
+        *,
+        betas: Sequence[float],
+        model_mean_type: str,
+        model_var_type: str,
+        channel_scales: Optional[np.ndarray] = None,
+        channel_biases: Optional[np.ndarray] = None,
+    ):
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.channel_scales = channel_scales
+        self.channel_biases = channel_biases
+
+        # Use float64 for accuracy
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.num_timesteps = int(betas.shape[0])
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        )
+
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+            (1.0 - self.alphas_cumprod_prev)
+            * np.sqrt(alphas)
+            / (1.0 - self.alphas_cumprod)
+        )
+
+    def scale_channels(self, x: th.Tensor) -> th.Tensor:
+        """Apply channel-wise scaling."""
+        if self.channel_scales is not None:
+            x = x * th.from_numpy(self.channel_scales).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        if self.channel_biases is not None:
+            x = x + th.from_numpy(self.channel_biases).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        return x
+
+    def unscale_channels(self, x: th.Tensor) -> th.Tensor:
+        """Remove channel-wise scaling."""
+        if self.channel_biases is not None:
+            x = x - th.from_numpy(self.channel_biases).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        if self.channel_scales is not None:
+            x = x / th.from_numpy(self.channel_scales).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        return x
+
+    def unscale_out_dict(
+        self, out: Dict[str, Union[th.Tensor, Any]]
+    ) -> Dict[str, Union[th.Tensor, Any]]:
+        return {
+            k: (self.unscale_channels(v) if isinstance(v, th.Tensor) else v)
+            for k, v in out.items()
+        }
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+
+            q(x_{t-1} | x_t, x_0)
+
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        assert (
+            posterior_mean.shape[0]
+            == posterior_variance.shape[0]
+            == posterior_log_variance_clipped.shape[0]
+            == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(
+        self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None
+    ):
+        """
+        Apply the model to get p(x_{t-1} | x_t).
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+
+        # Direct prediction of eps
+        model_output = model(x, t, **model_kwargs)
+        if isinstance(model_output, tuple):
+            model_output, prev_latent = model_output
+            model_kwargs["prev_latent"] = prev_latent
+
+        # Convert model output to mean and variance
+        model_variance, model_log_variance = {
+            # for fixedlarge, we set the initial (log-)variance like so
+            # to get a better decoder log likelihood.
+            "fixed_large": (
+                np.append(self.posterior_variance[1], self.betas[1:]),
+                np.log(np.append(self.posterior_variance[1], self.betas[1:])),
+            ),
+            "fixed_small": (
+                self.posterior_variance,
+                self.posterior_log_variance_clipped,
+            ),
+        }[self.model_var_type]
+        model_variance = _extract_into_tensor(model_variance, t, x.shape)
+        model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
+
+        def process_xstart(x):
+            if denoised_fn is not None:
+                x = denoised_fn(x)
+            if clip_denoised:
+                x = x.clamp(
+                    -self.channel_scales[0] * 0.67, self.channel_scales[0] * 0.67
+                )
+                x[:, 3:] = x[:, 3:].clamp(
+                    -self.channel_scales[3] * 0.5, self.channel_scales[3] * 0.5
+                )
+                return x
+            return x
+
+        if self.model_mean_type == "x_prev":
+            pred_xstart = process_xstart(
+                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
+            )
+            model_mean = model_output
+        elif self.model_mean_type in ["x_start", "epsilon"]:
+            if self.model_mean_type == "x_start":
+                pred_xstart = process_xstart(model_output)
+            else:
+                pred_xstart = process_xstart(
+                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+                )
+            model_mean, _, _ = self.q_posterior_mean_variance(
+                x_start=pred_xstart, x_t=x, t=t
+            )
+            # print('p_mean_variance:', pred_xstart.min(), pred_xstart.max())
+        else:
+            raise NotImplementedError(self.model_mean_type)
+
+        assert (
+            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+        )
+        return {
+            "mean": model_mean,
+            "variance": model_variance,
+            "log_variance": model_log_variance,
+            "pred_xstart": pred_xstart,
+        }
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def _predict_xstart_from_xprev(self, x_t, t, xprev):
+        assert x_t.shape == xprev.shape
+        return (  # (xprev - coef2*x_t) / coef1
+            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
+            - _extract_into_tensor(
+                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
+            )
+            * x_t
+        )
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def ddim_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples.
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield self.unscale_out_dict(out)
+                img = out["sample"]
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+        sigma = (
+            eta
+            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        )
+
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma**2) * eps
+        )
+        nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    A diffusion process which can skip steps in a base diffusion process.
+    """
+
+    def __init__(self, use_timesteps: Iterable[int], **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)
+        last_alpha_cumprod = 1.0
+        new_betas = []
+        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+            if i in self.use_timesteps:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(self, model, *args, **kwargs):
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def _wrap_model(self, model):
+        if isinstance(model, _WrappedModel):
+            return model
+        return _WrappedModel(model, self.timestep_map, self.original_num_steps)
+
+
+class _WrappedModel:
+    """Helper class to wrap models for SpacedDiffusion."""
+
+    def __init__(self, model, timestep_map, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        return self.model(x, new_ts, **kwargs)

spar3d/models/diffusion/sampler.py

@@ -0,0 +1,134 @@
+# --------------------------------------------------------
+# Adapted from: https://github.com/openai/point-e
+# Licensed under the MIT License
+# Copyright (c) 2022 OpenAI
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# --------------------------------------------------------
+
+from typing import Dict, Iterator
+
+import torch
+import torch.nn as nn
+
+from .gaussian_diffusion import GaussianDiffusion
+
+
+class PointCloudSampler:
+    """
+    A wrapper around a model that produces conditional sample tensors.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        diffusion: GaussianDiffusion,
+        num_points: int,
+        point_dim: int = 3,
+        guidance_scale: float = 3.0,
+        clip_denoised: bool = True,
+        sigma_min: float = 1e-3,
+        sigma_max: float = 120,
+        s_churn: float = 3,
+    ):
+        self.model = model
+        self.num_points = num_points
+        self.point_dim = point_dim
+        self.guidance_scale = guidance_scale
+        self.clip_denoised = clip_denoised
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.s_churn = s_churn
+
+        self.diffusion = diffusion
+
+    def sample_batch_progressive(
+        self,
+        batch_size: int,
+        condition: torch.Tensor,
+        noise=None,
+        device=None,
+        guidance_scale=None,
+    ) -> Iterator[Dict[str, torch.Tensor]]:
+        """
+        Generate samples progressively using classifier-free guidance.
+
+        Args:
+            batch_size: Number of samples to generate
+            condition: Conditioning tensor
+            noise: Optional initial noise tensor
+            device: Device to run on
+            guidance_scale: Optional override for guidance scale
+
+        Returns:
+            Iterator of dicts containing intermediate samples
+        """
+        if guidance_scale is None:
+            guidance_scale = self.guidance_scale
+
+        sample_shape = (batch_size, self.point_dim, self.num_points)
+
+        # Double the batch for classifier-free guidance
+        if guidance_scale != 1 and guidance_scale != 0:
+            condition = torch.cat([condition, torch.zeros_like(condition)], dim=0)
+            if noise is not None:
+                noise = torch.cat([noise, noise], dim=0)
+        model_kwargs = {"condition": condition}
+
+        internal_batch_size = batch_size
+        if guidance_scale != 1 and guidance_scale != 0:
+            model = self._uncond_guide_model(self.model, guidance_scale)
+            internal_batch_size *= 2
+        else:
+            model = self.model
+
+        samples_it = self.diffusion.ddim_sample_loop_progressive(
+            model,
+            shape=(internal_batch_size, *sample_shape[1:]),
+            model_kwargs=model_kwargs,
+            device=device,
+            clip_denoised=self.clip_denoised,
+            noise=noise,
+        )
+
+        for x in samples_it:
+            samples = {
+                "xstart": x["pred_xstart"][:batch_size],
+                "xprev": x["sample"][:batch_size] if "sample" in x else x["x"],
+            }
+            yield samples
+
+    def _uncond_guide_model(self, model: nn.Module, scale: float) -> nn.Module:
+        """
+        Wraps the model for classifier-free guidance.
+        """
+
+        def model_fn(x_t, ts, **kwargs):
+            half = x_t[: len(x_t) // 2]
+            combined = torch.cat([half, half], dim=0)
+            model_out = model(combined, ts, **kwargs)
+
+            eps, rest = model_out[:, : self.point_dim], model_out[:, self.point_dim :]
+            cond_eps, uncond_eps = torch.chunk(eps, 2, dim=0)
+            half_eps = uncond_eps + scale * (cond_eps - uncond_eps)
+            eps = torch.cat([half_eps, half_eps], dim=0)
+            return torch.cat([eps, rest], dim=1)
+
+        return model_fn

spar3d/models/global_estimator/reni_estimator.py

@@ -0,0 +1,112 @@
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from jaxtyping import Float
+from torch import Tensor
+
+from spar3d.models.illumination.reni.env_map import RENIEnvMap
+from spar3d.models.utils import BaseModule
+
+
+def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
+    assert d6.shape[-1] == 6, "Input tensor must have shape (..., 6)"
+
+    def proj_u2a(u, a):
+        r"""
+        u: batch x 3
+        a: batch x 3
+        """
+        inner_prod = torch.sum(u * a, dim=-1, keepdim=True)
+        norm2 = torch.sum(u**2, dim=-1, keepdim=True)
+        norm2 = torch.clamp(norm2, min=1e-8)
+        factor = inner_prod / (norm2 + 1e-10)
+        return factor * u
+
+    x_raw, y_raw = d6[..., :3], d6[..., 3:]
+
+    x = F.normalize(x_raw, dim=-1)
+    y = F.normalize(y_raw - proj_u2a(x, y_raw), dim=-1)
+    z = torch.cross(x, y, dim=-1)
+
+    return torch.stack((x, y, z), dim=-1)
+
+
+class ReniLatentCodeEstimator(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        triplane_features: int = 40
+
+        n_layers: int = 5
+        hidden_features: int = 512
+        activation: str = "relu"
+
+        pool: str = "mean"
+
+        reni_env_config: dict = field(default_factory=dict)
+
+    cfg: Config
+
+    def configure(self):
+        layers = []
+        cur_features = self.cfg.triplane_features * 3
+        for _ in range(self.cfg.n_layers):
+            layers.append(
+                nn.Conv2d(
+                    cur_features,
+                    self.cfg.hidden_features,
+                    kernel_size=3,
+                    padding=0,
+                    stride=2,
+                )
+            )
+            layers.append(self.make_activation(self.cfg.activation))
+
+            cur_features = self.cfg.hidden_features
+
+        self.layers = nn.Sequential(*layers)
+
+        self.reni_env_map = RENIEnvMap(self.cfg.reni_env_config)
+        self.latent_dim = self.reni_env_map.field.latent_dim
+
+        self.fc_latents = nn.Linear(self.cfg.hidden_features, self.latent_dim * 3)
+        nn.init.normal_(self.fc_latents.weight, mean=0.0, std=0.3)
+
+        self.fc_rotations = nn.Linear(self.cfg.hidden_features, 6)
+        nn.init.constant_(self.fc_rotations.bias, 0.0)
+        nn.init.normal_(
+            self.fc_rotations.weight, mean=0.0, std=0.01
+        )  # Small variance here
+
+        self.fc_scale = nn.Linear(self.cfg.hidden_features, 1)
+        nn.init.constant_(self.fc_scale.bias, 0.0)
+        nn.init.normal_(self.fc_scale.weight, mean=0.0, std=0.01)  # Small variance here
+
+    def make_activation(self, activation):
+        if activation == "relu":
+            return nn.ReLU(inplace=True)
+        elif activation == "silu":
+            return nn.SiLU(inplace=True)
+        else:
+            raise NotImplementedError
+
+    def forward(
+        self,
+        triplane: Float[Tensor, "B 3 F Ht Wt"],
+    ) -> dict[str, Any]:
+        x = self.layers(
+            triplane.reshape(
+                triplane.shape[0], -1, triplane.shape[-2], triplane.shape[-1]
+            )
+        )
+        x = x.mean(dim=[-2, -1])
+
+        latents = self.fc_latents(x).reshape(-1, self.latent_dim, 3)
+        rotations = self.fc_rotations(x)
+        scale = self.fc_scale(x)
+
+        env_map = self.reni_env_map(latents, rotation_6d_to_matrix(rotations), scale)
+
+        return {"illumination": env_map["rgb"]}

spar3d/models/illumination/reni/components/film_siren.py

@@ -0,0 +1,148 @@
+"""FiLM Siren MLP as per https://marcoamonteiro.github.io/pi-GAN-website/."""
+
+from typing import Optional
+
+import numpy as np
+import torch
+from torch import nn
+
+
+def kaiming_leaky_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Linear") != -1:
+        torch.nn.init.kaiming_normal_(
+            m.weight, a=0.2, mode="fan_in", nonlinearity="leaky_relu"
+        )
+
+
+def frequency_init(freq):
+    def init(m):
+        with torch.no_grad():
+            if isinstance(m, nn.Linear):
+                num_input = m.weight.size(-1)
+                m.weight.uniform_(
+                    -np.sqrt(6 / num_input) / freq, np.sqrt(6 / num_input) / freq
+                )
+
+    return init
+
+
+def first_layer_film_sine_init(m):
+    with torch.no_grad():
+        if isinstance(m, nn.Linear):
+            num_input = m.weight.size(-1)
+            m.weight.uniform_(-1 / num_input, 1 / num_input)
+
+
+class CustomMappingNetwork(nn.Module):
+    def __init__(self, in_features, map_hidden_layers, map_hidden_dim, map_output_dim):
+        super().__init__()
+
+        self.network = []
+
+        for _ in range(map_hidden_layers):
+            self.network.append(nn.Linear(in_features, map_hidden_dim))
+            self.network.append(nn.LeakyReLU(0.2, inplace=True))
+            in_features = map_hidden_dim
+
+        self.network.append(nn.Linear(map_hidden_dim, map_output_dim))
+
+        self.network = nn.Sequential(*self.network)
+
+        self.network.apply(kaiming_leaky_init)
+        with torch.no_grad():
+            self.network[-1].weight *= 0.25
+
+    def forward(self, z):
+        frequencies_offsets = self.network(z)
+        frequencies = frequencies_offsets[
+            ..., : torch.div(frequencies_offsets.shape[-1], 2, rounding_mode="floor")
+        ]
+        phase_shifts = frequencies_offsets[
+            ..., torch.div(frequencies_offsets.shape[-1], 2, rounding_mode="floor") :
+        ]
+
+        return frequencies, phase_shifts
+
+
+class FiLMLayer(nn.Module):
+    def __init__(self, input_dim, hidden_dim):
+        super().__init__()
+        self.layer = nn.Linear(input_dim, hidden_dim)
+
+    def forward(self, x, freq, phase_shift):
+        x = self.layer(x)
+        freq = freq.expand_as(x)
+        phase_shift = phase_shift.expand_as(x)
+        return torch.sin(freq * x + phase_shift)
+
+
+class FiLMSiren(nn.Module):
+    """FiLM Conditioned Siren network."""
+
+    def __init__(
+        self,
+        in_dim: int,
+        hidden_layers: int,
+        hidden_features: int,
+        mapping_network_in_dim: int,
+        mapping_network_layers: int,
+        mapping_network_features: int,
+        out_dim: int,
+        outermost_linear: bool = False,
+        out_activation: Optional[nn.Module] = None,
+    ) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        assert self.in_dim > 0
+        self.out_dim = out_dim if out_dim is not None else hidden_features
+        self.hidden_layers = hidden_layers
+        self.hidden_features = hidden_features
+        self.mapping_network_in_dim = mapping_network_in_dim
+        self.mapping_network_layers = mapping_network_layers
+        self.mapping_network_features = mapping_network_features
+        self.outermost_linear = outermost_linear
+        self.out_activation = out_activation
+
+        self.net = nn.ModuleList()
+
+        self.net.append(FiLMLayer(self.in_dim, self.hidden_features))
+
+        for _ in range(self.hidden_layers - 1):
+            self.net.append(FiLMLayer(self.hidden_features, self.hidden_features))
+
+        self.final_layer = None
+        if self.outermost_linear:
+            self.final_layer = nn.Linear(self.hidden_features, self.out_dim)
+            self.final_layer.apply(frequency_init(25))
+        else:
+            final_layer = FiLMLayer(self.hidden_features, self.out_dim)
+            self.net.append(final_layer)
+
+        self.mapping_network = CustomMappingNetwork(
+            in_features=self.mapping_network_in_dim,
+            map_hidden_layers=self.mapping_network_layers,
+            map_hidden_dim=self.mapping_network_features,
+            map_output_dim=(len(self.net)) * self.hidden_features * 2,
+        )
+
+        self.net.apply(frequency_init(25))
+        self.net[0].apply(first_layer_film_sine_init)
+
+    def forward_with_frequencies_phase_shifts(self, x, frequencies, phase_shifts):
+        """Get conditiional frequencies and phase shifts from mapping network."""
+        frequencies = frequencies * 15 + 30
+
+        for index, layer in enumerate(self.net):
+            start = index * self.hidden_features
+            end = (index + 1) * self.hidden_features
+            x = layer(x, frequencies[..., start:end], phase_shifts[..., start:end])
+
+        x = self.final_layer(x) if self.final_layer is not None else x
+        output = self.out_activation(x) if self.out_activation is not None else x
+        return output
+
+    def forward(self, x, conditioning_input):
+        """Forward pass."""
+        frequencies, phase_shifts = self.mapping_network(conditioning_input)
+        return self.forward_with_frequencies_phase_shifts(x, frequencies, phase_shifts)

spar3d/models/illumination/reni/components/siren.py

@@ -0,0 +1,118 @@
+"""Siren MLP https://www.vincentsitzmann.com/siren/"""
+
+from typing import Optional
+
+import numpy as np
+import torch
+from torch import nn
+
+
+class SineLayer(nn.Module):
+    """
+    Sine layer for the SIREN network.
+    """
+
+    def __init__(
+        self, in_features, out_features, bias=True, is_first=False, omega_0=30.0
+    ):
+        super().__init__()
+        self.omega_0 = omega_0
+        self.is_first = is_first
+
+        self.in_features = in_features
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+
+        self.init_weights()
+
+    def init_weights(self):
+        with torch.no_grad():
+            if self.is_first:
+                self.linear.weight.uniform_(-1 / self.in_features, 1 / self.in_features)
+            else:
+                self.linear.weight.uniform_(
+                    -np.sqrt(6 / self.in_features) / self.omega_0,
+                    np.sqrt(6 / self.in_features) / self.omega_0,
+                )
+
+    def forward(self, x):
+        return torch.sin(self.omega_0 * self.linear(x))
+
+
+class Siren(nn.Module):
+    """Siren network.
+
+    Args:
+        in_dim: Input layer dimension
+        num_layers: Number of network layers
+        layer_width: Width of each MLP layer
+        out_dim: Output layer dimension. Uses layer_width if None.
+        activation: intermediate layer activation function.
+        out_activation: output activation function.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        hidden_layers: int,
+        hidden_features: int,
+        out_dim: Optional[int] = None,
+        outermost_linear: bool = False,
+        first_omega_0: float = 30,
+        hidden_omega_0: float = 30,
+        out_activation: Optional[nn.Module] = None,
+    ) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        assert self.in_dim > 0
+        self.out_dim = out_dim if out_dim is not None else hidden_features
+        self.outermost_linear = outermost_linear
+        self.first_omega_0 = first_omega_0
+        self.hidden_omega_0 = hidden_omega_0
+        self.hidden_layers = hidden_layers
+        self.layer_width = hidden_features
+        self.out_activation = out_activation
+
+        self.net = []
+        self.net.append(
+            SineLayer(in_dim, hidden_features, is_first=True, omega_0=first_omega_0)
+        )
+
+        for _ in range(hidden_layers):
+            self.net.append(
+                SineLayer(
+                    hidden_features,
+                    hidden_features,
+                    is_first=False,
+                    omega_0=hidden_omega_0,
+                )
+            )
+
+        if outermost_linear:
+            final_layer = nn.Linear(hidden_features, self.out_dim)
+
+            with torch.no_grad():
+                final_layer.weight.uniform_(
+                    -np.sqrt(6 / hidden_features) / hidden_omega_0,
+                    np.sqrt(6 / hidden_features) / hidden_omega_0,
+                )
+
+            self.net.append(final_layer)
+        else:
+            self.net.append(
+                SineLayer(
+                    hidden_features,
+                    self.out_dim,
+                    is_first=False,
+                    omega_0=hidden_omega_0,
+                )
+            )
+
+        if self.out_activation is not None:
+            self.net.append(self.out_activation)
+
+        self.net = nn.Sequential(*self.net)
+
+    def forward(self, model_input):
+        """Forward pass through the network"""
+        output = self.net(model_input)
+        return output

spar3d/models/illumination/reni/components/transformer_decoder.py

@@ -0,0 +1,189 @@
+from typing import Optional
+
+import torch
+from torch import nn
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        direction_input_dim: int,
+        conditioning_input_dim: int,
+        latent_dim: int,
+        num_heads: int,
+    ):
+        """
+        Multi-Head Attention module.
+
+        Args:
+            direction_input_dim (int): The input dimension of the directional input.
+            conditioning_input_dim (int): The input dimension of the conditioning input.
+            latent_dim (int): The latent dimension of the module.
+            num_heads (int): The number of heads to use in the attention mechanism.
+        """
+        super().__init__()
+        assert latent_dim % num_heads == 0, "latent_dim must be divisible by num_heads"
+        self.num_heads = num_heads
+        self.head_dim = latent_dim // num_heads
+        self.scale = self.head_dim**-0.5
+
+        self.query = nn.Linear(direction_input_dim, latent_dim)
+        self.key = nn.Linear(conditioning_input_dim, latent_dim)
+        self.value = nn.Linear(conditioning_input_dim, latent_dim)
+        self.fc_out = nn.Linear(latent_dim, latent_dim)
+
+    def forward(
+        self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Forward pass of the Multi-Head Attention module.
+
+        Args:
+            query (torch.Tensor): The directional input tensor.
+            key (torch.Tensor): The conditioning input tensor for the keys.
+            value (torch.Tensor): The conditioning input tensor for the values.
+
+        Returns:
+            torch.Tensor: The output tensor of the Multi-Head Attention module.
+        """
+        batch_size = query.size(0)
+
+        Q = (
+            self.query(query)
+            .view(batch_size, -1, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+        K = (
+            self.key(key)
+            .view(batch_size, -1, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+        V = (
+            self.value(value)
+            .view(batch_size, -1, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+
+        attention = (
+            torch.einsum("bnqk,bnkh->bnqh", [Q, K.transpose(-2, -1)]) * self.scale
+        )
+        attention = torch.softmax(attention, dim=-1)
+
+        out = torch.einsum("bnqh,bnhv->bnqv", [attention, V])
+        out = (
+            out.transpose(1, 2)
+            .contiguous()
+            .view(batch_size, -1, self.num_heads * self.head_dim)
+        )
+
+        out = self.fc_out(out).squeeze(1)
+        return out
+
+
+class AttentionLayer(nn.Module):
+    def __init__(
+        self,
+        direction_input_dim: int,
+        conditioning_input_dim: int,
+        latent_dim: int,
+        num_heads: int,
+    ):
+        """
+        Attention Layer module.
+
+        Args:
+            direction_input_dim (int): The input dimension of the directional input.
+            conditioning_input_dim (int): The input dimension of the conditioning input.
+            latent_dim (int): The latent dimension of the module.
+            num_heads (int): The number of heads to use in the attention mechanism.
+        """
+        super().__init__()
+        self.mha = MultiHeadAttention(
+            direction_input_dim, conditioning_input_dim, latent_dim, num_heads
+        )
+        self.norm1 = nn.LayerNorm(latent_dim)
+        self.norm2 = nn.LayerNorm(latent_dim)
+        self.fc = nn.Sequential(
+            nn.Linear(latent_dim, latent_dim),
+            nn.ReLU(),
+            nn.Linear(latent_dim, latent_dim),
+        )
+
+    def forward(
+        self, directional_input: torch.Tensor, conditioning_input: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Forward pass of the Attention Layer module.
+
+        Args:
+            directional_input (torch.Tensor): The directional input tensor.
+            conditioning_input (torch.Tensor): The conditioning input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor of the Attention Layer module.
+        """
+        attn_output = self.mha(
+            directional_input, conditioning_input, conditioning_input
+        )
+        out1 = self.norm1(attn_output + directional_input)
+        fc_output = self.fc(out1)
+        out2 = self.norm2(fc_output + out1)
+        return out2
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        conditioning_input_dim: int,
+        hidden_features: int,
+        num_heads: int,
+        num_layers: int,
+        out_activation: Optional[nn.Module],
+    ):
+        """
+        Decoder module.
+
+        Args:
+            in_dim (int): The input dimension of the module.
+            conditioning_input_dim (int): The input dimension of the conditioning input.
+            hidden_features (int): The number of hidden features in the module.
+            num_heads (int): The number of heads to use in the attention mechanism.
+            num_layers (int): The number of layers in the module.
+            out_activation (nn.Module): The activation function to use on the output tensor.
+        """
+        super().__init__()
+        self.residual_projection = nn.Linear(
+            in_dim, hidden_features
+        )  # projection for residual connection
+        self.layers = nn.ModuleList(
+            [
+                AttentionLayer(
+                    hidden_features, conditioning_input_dim, hidden_features, num_heads
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.fc = nn.Linear(hidden_features, 3)  # 3 for RGB
+        self.out_activation = out_activation
+
+    def forward(
+        self, x: torch.Tensor, conditioning_input: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Forward pass of the Decoder module.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+            conditioning_input (torch.Tensor): The conditioning input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor of the Decoder module.
+        """
+        x = self.residual_projection(x)
+        for layer in self.layers:
+            x = layer(x, conditioning_input)
+        x = self.fc(x)
+        if self.out_activation is not None:
+            x = self.out_activation(x)
+        return x

spar3d/models/illumination/reni/components/vn_layers.py

@@ -0,0 +1,548 @@
+# MIT License
+
+# Copyright (c) 2022 Phil Wang
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""All code taken from https://github.com/lucidrains/VN-transformer"""
+
+from collections import namedtuple
+from functools import wraps
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, reduce
+from einops.layers.torch import Rearrange
+from packaging import version
+from torch import einsum, nn
+
+# constants
+
+FlashAttentionConfig = namedtuple(
+    "FlashAttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"]
+)
+
+# helpers
+
+
+def exists(val):
+    return val is not None
+
+
+def once(fn):
+    called = False
+
+    @wraps(fn)
+    def inner(x):
+        nonlocal called
+        if called:
+            return
+        called = True
+        return fn(x)
+
+    return inner
+
+
+print_once = once(print)
+
+# main class
+
+
+class Attend(nn.Module):
+    def __init__(self, dropout=0.0, flash=False, l2_dist=False):
+        super().__init__()
+        assert not (
+            flash and l2_dist
+        ), "flash attention is not compatible with l2 distance"
+        self.l2_dist = l2_dist
+
+        self.dropout = dropout
+        self.attn_dropout = nn.Dropout(dropout)
+
+        self.flash = flash
+        assert not (
+            flash and version.parse(torch.__version__) < version.parse("2.0.0")
+        ), "in order to use flash attention, you must be using pytorch 2.0 or above"
+
+        # determine efficient attention configs for cuda and cpu
+
+        self.cpu_config = FlashAttentionConfig(True, True, True)
+        self.cuda_config = None
+
+        if not torch.cuda.is_available() or not flash:
+            return
+
+        device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
+
+        if device_properties.major == 8 and device_properties.minor == 0:
+            print_once(
+                "A100 GPU detected, using flash attention if input tensor is on cuda"
+            )
+            self.cuda_config = FlashAttentionConfig(True, False, False)
+        else:
+            print_once(
+                "Non-A100 GPU detected, using math or mem efficient attention if input tensor is on cuda"
+            )
+            self.cuda_config = FlashAttentionConfig(False, True, True)
+
+    def flash_attn(self, q, k, v, mask=None):
+        _, heads, q_len, _, _, is_cuda = (
+            *q.shape,
+            k.shape[-2],
+            q.is_cuda,
+        )
+
+        # Check if mask exists and expand to compatible shape
+        # The mask is B L, so it would have to be expanded to B H N L
+
+        if exists(mask):
+            mask = mask.expand(-1, heads, q_len, -1)
+
+        # Check if there is a compatible device for flash attention
+
+        config = self.cuda_config if is_cuda else self.cpu_config
+
+        # pytorch 2.0 flash attn: q, k, v, mask, dropout, softmax_scale
+
+        with torch.backends.cuda.sdp_kernel(**config._asdict()):
+            out = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=mask,
+                dropout_p=self.dropout if self.training else 0.0,
+            )
+
+        return out
+
+    def forward(self, q, k, v, mask=None):
+        """
+        einstein notation
+        b - batch
+        h - heads
+        n, i, j - sequence length (base sequence length, source, target)
+        d - feature dimension
+        """
+        scale = q.shape[-1] ** -0.5
+
+        if exists(mask) and mask.ndim != 4:
+            mask = rearrange(mask, "b j -> b 1 1 j")
+
+        if self.flash:
+            return self.flash_attn(q, k, v, mask=mask)
+
+        # similarity
+
+        sim = einsum("b h i d, b h j d -> b h i j", q, k) * scale
+
+        # l2 distance
+
+        if self.l2_dist:
+            # -cdist squared == (-q^2 + 2qk - k^2)
+            # so simply work off the qk above
+            q_squared = reduce(q**2, "b h i d -> b h i 1", "sum")
+            k_squared = reduce(k**2, "b h j d -> b h 1 j", "sum")
+            sim = sim * 2 - q_squared - k_squared
+
+        # key padding mask
+
+        if exists(mask):
+            sim = sim.masked_fill(~mask, -torch.finfo(sim.dtype).max)
+
+        # attention
+
+        attn = sim.softmax(dim=-1)
+        attn = self.attn_dropout(attn)
+
+        # aggregate values
+
+        out = einsum("b h i j, b h j d -> b h i d", attn, v)
+
+        return out
+
+
+# helper
+
+
+def exists(val):  # noqa: F811
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def inner_dot_product(x, y, *, dim=-1, keepdim=True):
+    return (x * y).sum(dim=dim, keepdim=keepdim)
+
+
+# layernorm
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.ones(dim))
+        self.register_buffer("beta", torch.zeros(dim))
+
+    def forward(self, x):
+        return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)
+
+
+# equivariant modules
+
+
+class VNLinear(nn.Module):
+    def __init__(self, dim_in, dim_out, bias_epsilon=0.0):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(dim_out, dim_in))
+
+        self.bias = None
+        self.bias_epsilon = bias_epsilon
+
+        # in this paper, they propose going for quasi-equivariance with a small bias, controllable with epsilon, which they claim lead to better stability and results
+
+        if bias_epsilon > 0.0:
+            self.bias = nn.Parameter(torch.randn(dim_out))
+
+    def forward(self, x):
+        out = einsum("... i c, o i -> ... o c", x, self.weight)
+
+        if exists(self.bias):
+            bias = F.normalize(self.bias, dim=-1) * self.bias_epsilon
+            out = out + rearrange(bias, "... -> ... 1")
+
+        return out
+
+
+class VNReLU(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.W = nn.Parameter(torch.randn(dim, dim))
+        self.U = nn.Parameter(torch.randn(dim, dim))
+
+    def forward(self, x):
+        q = einsum("... i c, o i -> ... o c", x, self.W)
+        k = einsum("... i c, o i -> ... o c", x, self.U)
+
+        qk = inner_dot_product(q, k)
+
+        k_norm = k.norm(dim=-1, keepdim=True).clamp(min=self.eps)
+        q_projected_on_k = q - inner_dot_product(q, k / k_norm) * k
+
+        out = torch.where(qk >= 0.0, q, q_projected_on_k)
+
+        return out
+
+
+class VNAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_head=64,
+        heads=8,
+        dim_coor=3,
+        bias_epsilon=0.0,
+        l2_dist_attn=False,
+        flash=False,
+        num_latents=None,  # setting this would enable perceiver-like cross attention from latents to sequence, with the latents derived from VNWeightedPool
+    ):
+        super().__init__()
+        assert not (
+            l2_dist_attn and flash
+        ), "l2 distance attention is not compatible with flash attention"
+
+        self.scale = (dim_coor * dim_head) ** -0.5
+        dim_inner = dim_head * heads
+        self.heads = heads
+
+        self.to_q_input = None
+        if exists(num_latents):
+            self.to_q_input = VNWeightedPool(
+                dim, num_pooled_tokens=num_latents, squeeze_out_pooled_dim=False
+            )
+
+        self.to_q = VNLinear(dim, dim_inner, bias_epsilon=bias_epsilon)
+        self.to_k = VNLinear(dim, dim_inner, bias_epsilon=bias_epsilon)
+        self.to_v = VNLinear(dim, dim_inner, bias_epsilon=bias_epsilon)
+        self.to_out = VNLinear(dim_inner, dim, bias_epsilon=bias_epsilon)
+
+        if l2_dist_attn and not exists(num_latents):
+            # tied queries and keys for l2 distance attention, and not perceiver-like attention
+            self.to_k = self.to_q
+
+        self.attend = Attend(flash=flash, l2_dist=l2_dist_attn)
+
+    def forward(self, x, mask=None):
+        """
+        einstein notation
+        b - batch
+        n - sequence
+        h - heads
+        d - feature dimension (channels)
+        c - coordinate dimension (3 for 3d space)
+        i - source sequence dimension
+        j - target sequence dimension
+        """
+
+        c = x.shape[-1]
+
+        if exists(self.to_q_input):
+            q_input = self.to_q_input(x, mask=mask)
+        else:
+            q_input = x
+
+        q, k, v = self.to_q(q_input), self.to_k(x), self.to_v(x)
+        q, k, v = map(
+            lambda t: rearrange(t, "b n (h d) c -> b h n (d c)", h=self.heads),
+            (q, k, v),
+        )
+
+        out = self.attend(q, k, v, mask=mask)
+
+        out = rearrange(out, "b h n (d c) -> b n (h d) c", c=c)
+        return self.to_out(out)
+
+
+def VNFeedForward(dim, mult=4, bias_epsilon=0.0):
+    dim_inner = int(dim * mult)
+    return nn.Sequential(
+        VNLinear(dim, dim_inner, bias_epsilon=bias_epsilon),
+        VNReLU(dim_inner),
+        VNLinear(dim_inner, dim, bias_epsilon=bias_epsilon),
+    )
+
+
+class VNLayerNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.ln = LayerNorm(dim)
+
+    def forward(self, x):
+        norms = x.norm(dim=-1)
+        x = x / rearrange(norms.clamp(min=self.eps), "... -> ... 1")
+        ln_out = self.ln(norms)
+        return x * rearrange(ln_out, "... -> ... 1")
+
+
+class VNWeightedPool(nn.Module):
+    def __init__(
+        self, dim, dim_out=None, num_pooled_tokens=1, squeeze_out_pooled_dim=True
+    ):
+        super().__init__()
+        dim_out = default(dim_out, dim)
+        self.weight = nn.Parameter(torch.randn(num_pooled_tokens, dim, dim_out))
+        self.squeeze_out_pooled_dim = num_pooled_tokens == 1 and squeeze_out_pooled_dim
+
+    def forward(self, x, mask=None):
+        if exists(mask):
+            mask = rearrange(mask, "b n -> b n 1 1")
+            x = x.masked_fill(~mask, 0.0)
+            numer = reduce(x, "b n d c -> b d c", "sum")
+            denom = mask.sum(dim=1)
+            mean_pooled = numer / denom.clamp(min=1e-6)
+        else:
+            mean_pooled = reduce(x, "b n d c -> b d c", "mean")
+
+        out = einsum("b d c, m d e -> b m e c", mean_pooled, self.weight)
+
+        if not self.squeeze_out_pooled_dim:
+            return out
+
+        out = rearrange(out, "b 1 d c -> b d c")
+        return out
+
+
+# equivariant VN transformer encoder
+
+
+class VNTransformerEncoder(nn.Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        depth,
+        dim_head=64,
+        heads=8,
+        dim_coor=3,
+        ff_mult=4,
+        final_norm=False,
+        bias_epsilon=0.0,
+        l2_dist_attn=False,
+        flash_attn=False,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.dim_coor = dim_coor
+
+        self.layers = nn.ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        VNAttention(
+                            dim=dim,
+                            dim_head=dim_head,
+                            heads=heads,
+                            bias_epsilon=bias_epsilon,
+                            l2_dist_attn=l2_dist_attn,
+                            flash=flash_attn,
+                        ),
+                        VNLayerNorm(dim),
+                        VNFeedForward(dim=dim, mult=ff_mult, bias_epsilon=bias_epsilon),
+                        VNLayerNorm(dim),
+                    ]
+                )
+            )
+
+        self.norm = VNLayerNorm(dim) if final_norm else nn.Identity()
+
+    def forward(self, x, mask=None):
+        *_, d, c = x.shape
+
+        assert (
+            x.ndim == 4 and d == self.dim and c == self.dim_coor
+        ), "input needs to be in the shape of (batch, seq, dim ({self.dim}), coordinate dim ({self.dim_coor}))"
+
+        for attn, attn_post_ln, ff, ff_post_ln in self.layers:
+            x = attn_post_ln(attn(x, mask=mask)) + x
+            x = ff_post_ln(ff(x)) + x
+
+        return self.norm(x)
+
+
+# invariant layers
+
+
+class VNInvariant(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_coor=3,
+    ):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            VNLinear(dim, dim_coor), VNReLU(dim_coor), Rearrange("... d e -> ... e d")
+        )
+
+    def forward(self, x):
+        return einsum("b n d i, b n i o -> b n o", x, self.mlp(x))
+
+
+# main class
+
+
+class VNTransformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        num_tokens=None,
+        dim_feat=None,
+        dim_head=64,
+        heads=8,
+        dim_coor=3,
+        reduce_dim_out=True,
+        bias_epsilon=0.0,
+        l2_dist_attn=False,
+        flash_attn=False,
+        translation_equivariance=False,
+        translation_invariant=False,
+    ):
+        super().__init__()
+        self.token_emb = nn.Embedding(num_tokens, dim) if exists(num_tokens) else None
+
+        dim_feat = default(dim_feat, 0)
+        self.dim_feat = dim_feat
+        self.dim_coor_total = dim_coor + dim_feat
+
+        assert (int(translation_equivariance) + int(translation_invariant)) <= 1
+        self.translation_equivariance = translation_equivariance
+        self.translation_invariant = translation_invariant
+
+        self.vn_proj_in = nn.Sequential(
+            Rearrange("... c -> ... 1 c"), VNLinear(1, dim, bias_epsilon=bias_epsilon)
+        )
+
+        self.encoder = VNTransformerEncoder(
+            dim=dim,
+            depth=depth,
+            dim_head=dim_head,
+            heads=heads,
+            bias_epsilon=bias_epsilon,
+            dim_coor=self.dim_coor_total,
+            l2_dist_attn=l2_dist_attn,
+            flash_attn=flash_attn,
+        )
+
+        if reduce_dim_out:
+            self.vn_proj_out = nn.Sequential(
+                VNLayerNorm(dim),
+                VNLinear(dim, 1, bias_epsilon=bias_epsilon),
+                Rearrange("... 1 c -> ... c"),
+            )
+        else:
+            self.vn_proj_out = nn.Identity()
+
+    def forward(
+        self, coors, *, feats=None, mask=None, return_concatted_coors_and_feats=False
+    ):
+        if self.translation_equivariance or self.translation_invariant:
+            coors_mean = reduce(coors, "... c -> c", "mean")
+            coors = coors - coors_mean
+
+        x = coors  # [batch, num_points, 3]
+
+        if exists(feats):
+            if feats.dtype == torch.long:
+                assert exists(
+                    self.token_emb
+                ), "num_tokens must be given to the VNTransformer (to build the Embedding), if the features are to be given as indices"
+                feats = self.token_emb(feats)
+
+            assert (
+                feats.shape[-1] == self.dim_feat
+            ), f"dim_feat should be set to {feats.shape[-1]}"
+            x = torch.cat((x, feats), dim=-1)  # [batch, num_points, 3 + dim_feat]
+
+        assert x.shape[-1] == self.dim_coor_total
+
+        x = self.vn_proj_in(x)  # [batch, num_points, hidden_dim, 3 + dim_feat]
+        x = self.encoder(x, mask=mask)  # [batch, num_points, hidden_dim, 3 + dim_feat]
+        x = self.vn_proj_out(x)  # [batch, num_points, 3 + dim_feat]
+
+        coors_out, feats_out = (
+            x[..., :3],
+            x[..., 3:],
+        )  # [batch, num_points, 3], [batch, num_points, dim_feat]
+
+        if self.translation_equivariance:
+            coors_out = coors_out + coors_mean
+
+        if not exists(feats):
+            return coors_out
+
+        if return_concatted_coors_and_feats:
+            return torch.cat((coors_out, feats_out), dim=-1)
+
+        return coors_out, feats_out

spar3d/models/illumination/reni/env_map.py

@@ -0,0 +1,93 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import torch
+from jaxtyping import Float
+from torch import Tensor
+
+from spar3d.models.utils import BaseModule
+
+from .field import RENIField
+
+
+def _direction_from_coordinate(
+    coordinate: Float[Tensor, "*B 2"],
+) -> Float[Tensor, "*B 3"]:
+    # OpenGL Convention
+    # +X Right
+    # +Y Up
+    # +Z Backward
+
+    u, v = coordinate.unbind(-1)
+    theta = (2 * torch.pi * u) - torch.pi
+    phi = torch.pi * v
+
+    dir = torch.stack(
+        [
+            theta.sin() * phi.sin(),
+            phi.cos(),
+            -1 * theta.cos() * phi.sin(),
+        ],
+        -1,
+    )
+    return dir
+
+
+def _get_sample_coordinates(
+    resolution: List[int], device: Optional[torch.device] = None
+) -> Float[Tensor, "H W 2"]:
+    return torch.stack(
+        torch.meshgrid(
+            (torch.arange(resolution[1], device=device) + 0.5) / resolution[1],
+            (torch.arange(resolution[0], device=device) + 0.5) / resolution[0],
+            indexing="xy",
+        ),
+        -1,
+    )
+
+
+class RENIEnvMap(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        reni_config: dict = field(default_factory=dict)
+        resolution: int = 128
+
+    cfg: Config
+
+    def configure(self):
+        self.field = RENIField(self.cfg.reni_config)
+        resolution = (self.cfg.resolution, self.cfg.resolution * 2)
+        sample_directions = _direction_from_coordinate(
+            _get_sample_coordinates(resolution)
+        )
+        self.img_shape = sample_directions.shape[:-1]
+
+        sample_directions_flat = sample_directions.view(-1, 3)
+        # Lastly these have y up but reni expects z up. Rotate 90 degrees on x axis
+        sample_directions_flat = torch.stack(
+            [
+                sample_directions_flat[:, 0],
+                -sample_directions_flat[:, 2],
+                sample_directions_flat[:, 1],
+            ],
+            -1,
+        )
+        self.sample_directions = torch.nn.Parameter(
+            sample_directions_flat, requires_grad=False
+        )
+
+    def forward(
+        self,
+        latent_codes: Float[Tensor, "B latent_dim 3"],
+        rotation: Optional[Float[Tensor, "B 3 3"]] = None,
+        scale: Optional[Float[Tensor, "B"]] = None,
+    ) -> Dict[str, Tensor]:
+        return {
+            k: v.view(latent_codes.shape[0], *self.img_shape, -1)
+            for k, v in self.field(
+                self.sample_directions.unsqueeze(0).repeat(latent_codes.shape[0], 1, 1),
+                latent_codes,
+                rotation=rotation,
+                scale=scale,
+            ).items()
+        }

spar3d/models/illumination/reni/field.py

@@ -0,0 +1,736 @@
+# Copyright 2023 The University of York. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Modified by Mark Boss
+
+"""RENI field"""
+
+import contextlib
+from dataclasses import dataclass
+from typing import Dict, Literal, Optional
+
+import torch
+from einops.layers.torch import Rearrange
+from jaxtyping import Float
+from torch import Tensor, nn
+
+from spar3d.models.network import get_activation_module, trunc_exp
+from spar3d.models.utils import BaseModule
+
+from .components.film_siren import FiLMSiren
+from .components.siren import Siren
+from .components.transformer_decoder import Decoder
+from .components.vn_layers import VNInvariant, VNLinear
+
+# from nerfstudio.cameras.rays import RaySamples
+
+
+def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:
+    """Computes the expected value of sin(y) where y ~ N(x_means, x_vars)
+
+    Args:
+        x_means: Mean values.
+        x_vars: Variance of values.
+
+    Returns:
+        torch.Tensor: The expected value of sin.
+    """
+
+    return torch.exp(-0.5 * x_vars) * torch.sin(x_means)
+
+
+class NeRFEncoding(torch.nn.Module):
+    """Multi-scale sinousoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
+    Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.
+
+    Args:
+        in_dim: Input dimension of tensor
+        num_frequencies: Number of encoded frequencies per axis
+        min_freq_exp: Minimum frequency exponent
+        max_freq_exp: Maximum frequency exponent
+        include_input: Append the input coordinate to the encoding
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        num_frequencies: int,
+        min_freq_exp: float,
+        max_freq_exp: float,
+        include_input: bool = False,
+        off_axis: bool = False,
+    ) -> None:
+        super().__init__()
+
+        self.in_dim = in_dim
+        self.num_frequencies = num_frequencies
+        self.min_freq = min_freq_exp
+        self.max_freq = max_freq_exp
+        self.include_input = include_input
+
+        self.off_axis = off_axis
+
+        self.P = torch.tensor(
+            [
+                [0.8506508, 0, 0.5257311],
+                [0.809017, 0.5, 0.309017],
+                [0.5257311, 0.8506508, 0],
+                [1, 0, 0],
+                [0.809017, 0.5, -0.309017],
+                [0.8506508, 0, -0.5257311],
+                [0.309017, 0.809017, -0.5],
+                [0, 0.5257311, -0.8506508],
+                [0.5, 0.309017, -0.809017],
+                [0, 1, 0],
+                [-0.5257311, 0.8506508, 0],
+                [-0.309017, 0.809017, -0.5],
+                [0, 0.5257311, 0.8506508],
+                [-0.309017, 0.809017, 0.5],
+                [0.309017, 0.809017, 0.5],
+                [0.5, 0.309017, 0.809017],
+                [0.5, -0.309017, 0.809017],
+                [0, 0, 1],
+                [-0.5, 0.309017, 0.809017],
+                [-0.809017, 0.5, 0.309017],
+                [-0.809017, 0.5, -0.309017],
+            ]
+        ).T
+
+    def get_out_dim(self) -> int:
+        if self.in_dim is None:
+            raise ValueError("Input dimension has not been set")
+        out_dim = self.in_dim * self.num_frequencies * 2
+
+        if self.off_axis:
+            out_dim = self.P.shape[1] * self.num_frequencies * 2
+
+        if self.include_input:
+            out_dim += self.in_dim
+        return out_dim
+
+    def forward(
+        self,
+        in_tensor: Float[Tensor, "*b input_dim"],
+        covs: Optional[Float[Tensor, "*b input_dim input_dim"]] = None,
+    ) -> Float[Tensor, "*b output_dim"]:
+        """Calculates NeRF encoding. If covariances are provided the encodings will be integrated as proposed
+            in mip-NeRF.
+
+        Args:
+            in_tensor: For best performance, the input tensor should be between 0 and 1.
+            covs: Covariances of input points.
+        Returns:
+            Output values will be between -1 and 1
+        """
+        # TODO check scaling here but just comment it for now
+        # in_tensor = 2 * torch.pi * in_tensor  # scale to [0, 2pi]
+        freqs = 2 ** torch.linspace(
+            self.min_freq, self.max_freq, self.num_frequencies
+        ).to(in_tensor.device)
+        # freqs = 2 ** (
+        #    torch.sin(torch.linspace(self.min_freq, torch.pi / 2.0, self.num_frequencies)) * self.max_freq
+        # ).to(in_tensor.device)
+        # freqs = 2 ** (
+        #     torch.linspace(self.min_freq, 1.0, self.num_frequencies).to(in_tensor.device) ** 0.2 * self.max_freq
+        # )
+
+        if self.off_axis:
+            scaled_inputs = (
+                torch.matmul(in_tensor, self.P.to(in_tensor.device))[..., None] * freqs
+            )
+        else:
+            scaled_inputs = (
+                in_tensor[..., None] * freqs
+            )  # [..., "input_dim", "num_scales"]
+        scaled_inputs = scaled_inputs.view(
+            *scaled_inputs.shape[:-2], -1
+        )  # [..., "input_dim" * "num_scales"]
+
+        if covs is None:
+            encoded_inputs = torch.sin(
+                torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1)
+            )
+        else:
+            input_var = (
+                torch.diagonal(covs, dim1=-2, dim2=-1)[..., :, None]
+                * freqs[None, :] ** 2
+            )
+            input_var = input_var.reshape((*input_var.shape[:-2], -1))
+            encoded_inputs = expected_sin(
+                torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1),
+                torch.cat(2 * [input_var], dim=-1),
+            )
+
+        if self.include_input:
+            encoded_inputs = torch.cat([encoded_inputs, in_tensor], dim=-1)
+        return encoded_inputs
+
+
+class RENIField(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        """Configuration for model instantiation"""
+
+        fixed_decoder: bool = False
+        """Whether to fix the decoder weights"""
+        equivariance: str = "SO2"
+        """Type of equivariance to use: None, SO2, SO3"""
+        axis_of_invariance: str = "y"
+        """Which axis should SO2 equivariance be invariant to: x, y, z"""
+        invariant_function: str = "GramMatrix"
+        """Type of invariant function to use: GramMatrix, VN"""
+        conditioning: str = "Concat"
+        """Type of conditioning to use: FiLM, Concat, Attention"""
+        positional_encoding: str = "NeRF"
+        """Type of positional encoding to use. Currently only NeRF is supported"""
+        encoded_input: str = "Directions"
+        """Type of input to encode: None, Directions, Conditioning, Both"""
+        latent_dim: int = 36
+        """Dimensionality of latent code, N for a latent code size of (N x 3)"""
+        hidden_layers: int = 3
+        """Number of hidden layers"""
+        hidden_features: int = 128
+        """Number of hidden features"""
+        mapping_layers: int = 3
+        """Number of mapping layers"""
+        mapping_features: int = 128
+        """Number of mapping features"""
+        num_attention_heads: int = 8
+        """Number of attention heads"""
+        num_attention_layers: int = 3
+        """Number of attention layers"""
+        out_features: int = 3  # RGB
+        """Number of output features"""
+        last_layer_linear: bool = False
+        """Whether to use a linear layer as the last layer"""
+        output_activation: str = "exp"
+        """Activation function for output layer: sigmoid, tanh, relu, exp, None"""
+        first_omega_0: float = 30.0
+        """Omega_0 for first layer"""
+        hidden_omega_0: float = 30.0
+        """Omega_0 for hidden layers"""
+        fixed_decoder: bool = False
+        """Whether to fix the decoder weights"""
+        old_implementation: bool = False
+        """Whether to match implementation of old RENI, when using old checkpoints"""
+
+    cfg: Config
+
+    def configure(self):
+        self.equivariance = self.cfg.equivariance
+        self.conditioning = self.cfg.conditioning
+        self.latent_dim = self.cfg.latent_dim
+        self.hidden_layers = self.cfg.hidden_layers
+        self.hidden_features = self.cfg.hidden_features
+        self.mapping_layers = self.cfg.mapping_layers
+        self.mapping_features = self.cfg.mapping_features
+        self.out_features = self.cfg.out_features
+        self.last_layer_linear = self.cfg.last_layer_linear
+        self.output_activation = self.cfg.output_activation
+        self.first_omega_0 = self.cfg.first_omega_0
+        self.hidden_omega_0 = self.cfg.hidden_omega_0
+        self.old_implementation = self.cfg.old_implementation
+        self.axis_of_invariance = ["x", "y", "z"].index(self.cfg.axis_of_invariance)
+
+        self.fixed_decoder = self.cfg.fixed_decoder
+        if self.cfg.invariant_function == "GramMatrix":
+            self.invariant_function = self.gram_matrix_invariance
+        else:
+            self.vn_proj_in = nn.Sequential(
+                Rearrange("... c -> ... 1 c"),
+                VNLinear(dim_in=1, dim_out=1, bias_epsilon=0),
+            )
+            dim_coor = 2 if self.cfg.equivariance == "SO2" else 3
+            self.vn_invar = VNInvariant(dim=1, dim_coor=dim_coor)
+            self.invariant_function = self.vn_invariance
+
+        self.network = self.setup_network()
+
+        if self.fixed_decoder:
+            for param in self.network.parameters():
+                param.requires_grad = False
+
+            if self.cfg.invariant_function == "VN":
+                for param in self.vn_proj_in.parameters():
+                    param.requires_grad = False
+                for param in self.vn_invar.parameters():
+                    param.requires_grad = False
+
+    @contextlib.contextmanager
+    def hold_decoder_fixed(self):
+        """Context manager to fix the decoder weights
+
+        Example usage:
+        ```
+        with instance_of_RENIField.hold_decoder_fixed():
+            # do stuff
+        ```
+        """
+        prev_state_network = {
+            name: p.requires_grad for name, p in self.network.named_parameters()
+        }
+        for param in self.network.parameters():
+            param.requires_grad = False
+        if self.cfg.invariant_function == "VN":
+            prev_state_proj_in = {
+                k: p.requires_grad for k, p in self.vn_proj_in.named_parameters()
+            }
+            prev_state_invar = {
+                k: p.requires_grad for k, p in self.vn_invar.named_parameters()
+            }
+            for param in self.vn_proj_in.parameters():
+                param.requires_grad = False
+            for param in self.vn_invar.parameters():
+                param.requires_grad = False
+
+        prev_decoder_state = self.fixed_decoder
+        self.fixed_decoder = True
+        try:
+            yield
+        finally:
+            # Restore the previous requires_grad state
+            for name, param in self.network.named_parameters():
+                param.requires_grad = prev_state_network[name]
+            if self.cfg.invariant_function == "VN":
+                for name, param in self.vn_proj_in.named_parameters():
+                    param.requires_grad_(prev_state_proj_in[name])
+                for name, param in self.vn_invar.named_parameters():
+                    param.requires_grad_(prev_state_invar[name])
+            self.fixed_decoder = prev_decoder_state
+
+    def vn_invariance(
+        self,
+        Z: Float[Tensor, "B latent_dim 3"],
+        D: Float[Tensor, "B num_rays 3"],
+        equivariance: Literal["None", "SO2", "SO3"] = "SO2",
+        axis_of_invariance: int = 1,
+    ):
+        """Generates a batched invariant representation from latent code Z and direction coordinates D.
+
+        Args:
+            Z: [B, latent_dim, 3] - Latent code.
+            D: [B num_rays, 3] - Direction coordinates.
+            equivariance: The type of equivariance to use. Options are 'None', 'SO2', 'SO3'.
+            axis_of_invariance: The axis of rotation invariance. Should be 0 (x-axis), 1 (y-axis), or 2 (z-axis).
+
+        Returns:
+            Tuple[Tensor, Tensor]: directional_input, conditioning_input
+        """
+        assert 0 <= axis_of_invariance < 3, "axis_of_invariance should be 0, 1, or 2."
+        other_axes = [i for i in range(3) if i != axis_of_invariance]
+
+        B, latent_dim, _ = Z.shape
+        _, num_rays, _ = D.shape
+
+        if equivariance == "None":
+            # get inner product between latent code and direction coordinates
+            innerprod = torch.sum(
+                Z.unsqueeze(1) * D.unsqueeze(2), dim=-1
+            )  # [B, num_rays, latent_dim]
+            z_input = (
+                Z.flatten(start_dim=1).unsqueeze(1).expand(B, num_rays, latent_dim * 3)
+            )  # [B, num_rays, latent_dim * 3]
+            return innerprod, z_input
+
+        if equivariance == "SO2":
+            z_other = torch.stack(
+                (Z[..., other_axes[0]], Z[..., other_axes[1]]), -1
+            )  # [B, latent_dim, 2]
+            d_other = torch.stack(
+                (D[..., other_axes[0]], D[..., other_axes[1]]), -1
+            ).unsqueeze(2)  # [B, num_rays, 1, 2]
+            d_other = d_other.expand(
+                B, num_rays, latent_dim, 2
+            )  # [B, num_rays, latent_dim, 2]
+
+            z_other_emb = self.vn_proj_in(z_other)  # [B, latent_dim, 1, 2]
+            z_other_invar = self.vn_invar(z_other_emb)  # [B, latent_dim, 2]
+
+            # Get invariant component of Z along the axis of invariance
+            z_invar = Z[..., axis_of_invariance].unsqueeze(-1)  # [B, latent_dim, 1]
+
+            # Innerproduct between projection of Z and D on the plane orthogonal to the axis of invariance.
+            # This encodes the rotational information. This is rotation-equivariant to rotations of either Z
+            # or D and is invariant to rotations of both Z and D.
+            innerprod = (z_other.unsqueeze(1) * d_other).sum(
+                dim=-1
+            )  # [B, num_rays, latent_dim]
+
+            # Compute norm along the axes orthogonal to the axis of invariance
+            d_other_norm = torch.sqrt(
+                D[..., other_axes[0]] ** 2 + D[..., other_axes[1]] ** 2
+            ).unsqueeze(-1)  # [B num_rays, 1]
+
+            # Get invariant component of D along the axis of invariance
+            d_invar = D[..., axis_of_invariance].unsqueeze(-1)  # [B, num_rays, 1]
+
+            directional_input = torch.cat(
+                (innerprod, d_invar, d_other_norm), -1
+            )  # [B, num_rays, latent_dim + 2]
+            conditioning_input = (
+                torch.cat((z_other_invar, z_invar), dim=-1)
+                .flatten(1)
+                .unsqueeze(1)
+                .expand(B, num_rays, latent_dim * 3)
+            )  # [B, num_rays, latent_dim * 3]
+
+            return directional_input, conditioning_input
+
+        if equivariance == "SO3":
+            z = self.vn_proj_in(Z)  # [B, latent_dim, 1, 3]
+            z_invar = self.vn_invar(z)  # [B, latent_dim, 3]
+            conditioning_input = (
+                z_invar.flatten(1).unsqueeze(1).expand(B, num_rays, latent_dim)
+            )  # [B, num_rays, latent_dim * 3]
+            # D [B, num_rays, 3] -> [B, num_rays, 1, 3]
+            # Z [B, latent_dim, 3] -> [B, 1, latent_dim, 3]
+            innerprod = torch.sum(
+                Z.unsqueeze(1) * D.unsqueeze(2), dim=-1
+            )  # [B, num_rays, latent_dim]
+            return innerprod, conditioning_input
+
+    def gram_matrix_invariance(
+        self,
+        Z: Float[Tensor, "B latent_dim 3"],
+        D: Float[Tensor, "B num_rays 3"],
+        equivariance: Literal["None", "SO2", "SO3"] = "SO2",
+        axis_of_invariance: int = 1,
+    ):
+        """Generates an invariant representation from latent code Z and direction coordinates D.
+
+        Args:
+            Z (torch.Tensor): Latent code (B x latent_dim x 3)
+            D (torch.Tensor): Direction coordinates (B x num_rays x 3)
+            equivariance (str): Type of equivariance to use. Options are 'none', 'SO2', and 'SO3'
+            axis_of_invariance (int): The axis of rotation invariance. Should be 0 (x-axis), 1 (y-axis), or 2 (z-axis).
+                Default is 1 (y-axis).
+        Returns:
+            torch.Tensor: Invariant representation
+        """
+        assert 0 <= axis_of_invariance < 3, "axis_of_invariance should be 0, 1, or 2."
+        other_axes = [i for i in range(3) if i != axis_of_invariance]
+
+        B, latent_dim, _ = Z.shape
+        _, num_rays, _ = D.shape
+
+        if equivariance == "None":
+            # get inner product between latent code and direction coordinates
+            innerprod = torch.sum(
+                Z.unsqueeze(1) * D.unsqueeze(2), dim=-1
+            )  # [B, num_rays, latent_dim]
+            z_input = (
+                Z.flatten(start_dim=1).unsqueeze(1).expand(B, num_rays, latent_dim * 3)
+            )  # [B, num_rays, latent_dim * 3]
+            return innerprod, z_input
+
+        if equivariance == "SO2":
+            # Select components along axes orthogonal to the axis of invariance
+            z_other = torch.stack(
+                (Z[..., other_axes[0]], Z[..., other_axes[1]]), -1
+            )  # [B, latent_dim, 2]
+            d_other = torch.stack(
+                (D[..., other_axes[0]], D[..., other_axes[1]]), -1
+            ).unsqueeze(2)  # [B, num_rays, 1, 2]
+            d_other = d_other.expand(
+                B, num_rays, latent_dim, 2
+            )  # size becomes [B, num_rays, latent_dim, 2]
+
+            # Invariant representation of Z, gram matrix G=Z*Z' is size num_rays x latent_dim x latent_dim
+            G = torch.bmm(z_other, torch.transpose(z_other, 1, 2))
+
+            # Flatten G to be size B x latent_dim^2
+            z_other_invar = G.flatten(start_dim=1)
+
+            # Get invariant component of Z along the axis of invariance
+            z_invar = Z[..., axis_of_invariance]  # [B, latent_dim]
+
+            # Innerprod is size num_rays x latent_dim
+            innerprod = (z_other.unsqueeze(1) * d_other).sum(
+                dim=-1
+            )  # [B, num_rays, latent_dim]
+
+            # Compute norm along the axes orthogonal to the axis of invariance
+            d_other_norm = torch.sqrt(
+                D[..., other_axes[0]] ** 2 + D[..., other_axes[1]] ** 2
+            ).unsqueeze(-1)  # [B, num_rays, 1]
+
+            # Get invariant component of D along the axis of invariance
+            d_invar = D[..., axis_of_invariance].unsqueeze(-1)  # [B, num_rays, 1]
+
+            if not self.old_implementation:
+                directional_input = torch.cat(
+                    (innerprod, d_invar, d_other_norm), -1
+                )  # [B, num_rays, latent_dim + 2]
+                conditioning_input = (
+                    torch.cat((z_other_invar, z_invar), -1)
+                    .unsqueeze(1)
+                    .expand(B, num_rays, latent_dim * 3)
+                )  # [B, num_rays, latent_dim^2 + latent_dim]
+            else:
+                # this is matching the previous implementation of RENI, needed if using old checkpoints
+                z_other_invar = z_other_invar.unsqueeze(1).expand(B, num_rays, -1)
+                z_invar = z_invar.unsqueeze(1).expand(B, num_rays, -1)
+                return torch.cat(
+                    (innerprod, z_other_invar, d_other_norm, z_invar, d_invar), 1
+                )
+
+            return directional_input, conditioning_input
+
+        if equivariance == "SO3":
+            G = Z @ torch.transpose(Z, 1, 2)  # [B, latent_dim, latent_dim]
+            innerprod = torch.sum(
+                Z.unsqueeze(1) * D.unsqueeze(2), dim=-1
+            )  # [B, num_rays, latent_dim]
+            z_invar = (
+                G.flatten(start_dim=1).unsqueeze(1).expand(B, num_rays, -1)
+            )  # [B, num_rays, latent_dim^2]
+            return innerprod, z_invar
+
+    def setup_network(self):
+        """Sets up the network architecture"""
+        base_input_dims = {
+            "VN": {
+                "None": {
+                    "direction": self.latent_dim,
+                    "conditioning": self.latent_dim * 3,
+                },
+                "SO2": {
+                    "direction": self.latent_dim + 2,
+                    "conditioning": self.latent_dim * 3,
+                },
+                "SO3": {
+                    "direction": self.latent_dim,
+                    "conditioning": self.latent_dim * 3,
+                },
+            },
+            "GramMatrix": {
+                "None": {
+                    "direction": self.latent_dim,
+                    "conditioning": self.latent_dim * 3,
+                },
+                "SO2": {
+                    "direction": self.latent_dim + 2,
+                    "conditioning": self.latent_dim**2 + self.latent_dim,
+                },
+                "SO3": {
+                    "direction": self.latent_dim,
+                    "conditioning": self.latent_dim**2,
+                },
+            },
+        }
+
+        # Extract the necessary input dimensions
+        input_types = ["direction", "conditioning"]
+        input_dims = {
+            key: base_input_dims[self.cfg.invariant_function][self.cfg.equivariance][
+                key
+            ]
+            for key in input_types
+        }
+
+        # Helper function to create NeRF encoding
+        def create_nerf_encoding(in_dim):
+            return NeRFEncoding(
+                in_dim=in_dim,
+                num_frequencies=2,
+                min_freq_exp=0.0,
+                max_freq_exp=2.0,
+                include_input=True,
+            )
+
+        # Dictionary-based encoding setup
+        encoding_setup = {
+            "None": [],
+            "Conditioning": ["conditioning"],
+            "Directions": ["direction"],
+            "Both": ["direction", "conditioning"],
+        }
+
+        # Setting up the required encodings
+        for input_type in encoding_setup.get(self.cfg.encoded_input, []):
+            # create self.{input_type}_encoding and update input_dims
+            setattr(
+                self,
+                f"{input_type}_encoding",
+                create_nerf_encoding(input_dims[input_type]),
+            )
+            input_dims[input_type] = getattr(
+                self, f"{input_type}_encoding"
+            ).get_out_dim()
+
+        output_activation = get_activation_module(self.cfg.output_activation)
+
+        network = None
+        if self.conditioning == "Concat":
+            network = Siren(
+                in_dim=input_dims["direction"] + input_dims["conditioning"],
+                hidden_layers=self.hidden_layers,
+                hidden_features=self.hidden_features,
+                out_dim=self.out_features,
+                outermost_linear=self.last_layer_linear,
+                first_omega_0=self.first_omega_0,
+                hidden_omega_0=self.hidden_omega_0,
+                out_activation=output_activation,
+            )
+        elif self.conditioning == "FiLM":
+            network = FiLMSiren(
+                in_dim=input_dims["direction"],
+                hidden_layers=self.hidden_layers,
+                hidden_features=self.hidden_features,
+                mapping_network_in_dim=input_dims["conditioning"],
+                mapping_network_layers=self.mapping_layers,
+                mapping_network_features=self.mapping_features,
+                out_dim=self.out_features,
+                outermost_linear=True,
+                out_activation=output_activation,
+            )
+        elif self.conditioning == "Attention":
+            # transformer where K, V is from conditioning input and Q is from pos encoded directional input
+            network = Decoder(
+                in_dim=input_dims["direction"],
+                conditioning_input_dim=input_dims["conditioning"],
+                hidden_features=self.cfg.hidden_features,
+                num_heads=self.cfg.num_attention_heads,
+                num_layers=self.cfg.num_attention_layers,
+                out_activation=output_activation,
+            )
+        assert network is not None, "unknown conditioning type"
+        return network
+
+    def apply_positional_encoding(self, directional_input, conditioning_input):
+        # conditioning on just invariant directional input
+        if self.cfg.encoded_input == "Conditioning":
+            conditioning_input = self.conditioning_encoding(
+                conditioning_input
+            )  # [num_rays, embedding_dim]
+        elif self.cfg.encoded_input == "Directions":
+            directional_input = self.direction_encoding(
+                directional_input
+            )  # [num_rays, embedding_dim]
+        elif self.cfg.encoded_input == "Both":
+            directional_input = self.direction_encoding(directional_input)
+            conditioning_input = self.conditioning_encoding(conditioning_input)
+
+        return directional_input, conditioning_input
+
+    def get_outputs(
+        self,
+        rays_d: Float[Tensor, "batch num_rays 3"],  # type: ignore
+        latent_codes: Float[Tensor, "batch_size latent_dim 3"],  # type: ignore
+        rotation: Optional[Float[Tensor, "batch_size 3 3"]] = None,  # type: ignore
+        scale: Optional[Float[Tensor, "batch_size"]] = None,  # type: ignore
+    ) -> Dict[str, Tensor]:
+        """Returns the outputs of the field.
+
+        Args:
+            ray_samples: [batch_size num_rays 3]
+            latent_codes: [batch_size, latent_dim, 3]
+            rotation: [batch_size, 3, 3]
+            scale: [batch_size]
+        """
+        if rotation is not None:
+            if len(rotation.shape) == 3:  # [batch_size, 3, 3]
+                # Expand latent_codes to match [batch_size, latent_dim, 3]
+                latent_codes = torch.einsum(
+                    "bik,blk->bli",
+                    rotation,
+                    latent_codes,
+                )
+            else:
+                raise NotImplementedError(
+                    "Unsupported rotation shape. Expected [batch_size, 3, 3]."
+                )
+
+        B, num_rays, _ = rays_d.shape
+        _, latent_dim, _ = latent_codes.shape
+
+        if not self.old_implementation:
+            directional_input, conditioning_input = self.invariant_function(
+                latent_codes,
+                rays_d,
+                equivariance=self.equivariance,
+                axis_of_invariance=self.axis_of_invariance,
+            )  # [B, num_rays, 3]
+
+            if self.cfg.positional_encoding == "NeRF":
+                directional_input, conditioning_input = self.apply_positional_encoding(
+                    directional_input, conditioning_input
+                )
+
+            if self.conditioning == "Concat":
+                model_outputs = self.network(
+                    torch.cat((directional_input, conditioning_input), dim=-1).reshape(
+                        B * num_rays, -1
+                    )
+                ).view(B, num_rays, 3)  # returns -> [B num_rays, 3]
+            elif self.conditioning == "FiLM":
+                model_outputs = self.network(
+                    directional_input.reshape(B * num_rays, -1),
+                    conditioning_input.reshape(B * num_rays, -1),
+                ).view(B, num_rays, 3)  # returns -> [B num_rays, 3]
+            elif self.conditioning == "Attention":
+                model_outputs = self.network(
+                    directional_input.reshape(B * num_rays, -1),
+                    conditioning_input.reshape(B * num_rays, -1),
+                ).view(B, num_rays, 3)  # returns -> [B num_rays, 3]
+        else:
+            # in the old implementation directions were sampled with y-up not z-up so need to swap y and z in directions
+            directions = torch.stack(
+                (rays_d[..., 0], rays_d[..., 2], rays_d[..., 1]), -1
+            )
+            model_input = self.invariant_function(
+                latent_codes,
+                directions,
+                equivariance=self.equivariance,
+                axis_of_invariance=self.axis_of_invariance,
+            )  # [B, num_rays, 3]
+
+            model_outputs = self.network(model_input.view(B * num_rays, -1)).view(
+                B, num_rays, 3
+            )
+
+        outputs = {}
+
+        if scale is not None:
+            scale = trunc_exp(scale)  # [num_rays] exp to ensure positive
+            model_outputs = model_outputs * scale.view(-1, 1, 1)  # [num_rays, 3]
+
+        outputs["rgb"] = model_outputs
+
+        return outputs
+
+    def forward(
+        self,
+        rays_d: Float[Tensor, "batch num_rays 3"],  # type: ignore
+        latent_codes: Float[Tensor, "batch_size latent_dim 3"],  # type: ignore
+        rotation: Optional[Float[Tensor, "batch_size 3 3"]] = None,  # type: ignore
+        scale: Optional[Float[Tensor, "batch_size"]] = None,  # type: ignore
+    ) -> Dict[str, Tensor]:
+        """Evaluates spherical field for a given ray bundle and rotation.
+
+        Args:
+            ray_samples: [B num_rays 3]
+            latent_codes: [B, num_rays, latent_dim, 3]
+            rotation: [batch_size, 3, 3]
+            scale: [batch_size]
+
+        Returns:
+            Dict[str, Tensor]: A dictionary containing the outputs of the field.
+        """
+        return self.get_outputs(
+            rays_d=rays_d,
+            latent_codes=latent_codes,
+            rotation=rotation,
+            scale=scale,
+        )

spar3d/models/image_estimator/clip_based_estimator.py

@@ -0,0 +1,184 @@
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+import alpha_clip
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+from torchvision.transforms import Normalize
+
+from spar3d.models.network import get_activation
+from spar3d.models.utils import BaseModule
+
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)
+
+
+@dataclass
+class HeadSpec:
+    name: str
+    out_channels: int
+    n_hidden_layers: int
+    output_activation: Optional[str] = None
+    output_bias: float = 0.0
+    add_to_decoder_features: bool = False
+    shape: Optional[list[int]] = None
+    distribution_eval: str = "sample"
+
+
+class ClipBasedHeadEstimator(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        model: str = "ViT-L/14@336px"
+
+        distribution: str = "beta"
+
+        # ["mean", "mode", "sample", "sample_mean"]
+        distribution_eval: str = "mode"
+
+        activation: str = "relu"
+        hidden_features: int = 512
+        heads: List[HeadSpec] = field(default_factory=lambda: [])
+
+    cfg: Config
+
+    def configure(self):
+        self.model, _ = alpha_clip.load(
+            self.cfg.model,
+        )  # change to your own ckpt path
+        self.model.eval()
+
+        if not hasattr(self.model.visual, "input_resolution"):
+            self.img_size = 224
+        else:
+            self.img_size = self.model.visual.input_resolution
+            # Check if img_size is subscribable and pick the first element
+            if hasattr(self.img_size, "__getitem__"):
+                self.img_size = self.img_size[0]
+
+        # Do not add the weights in self.model to the optimizer
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        assert len(self.cfg.heads) > 0
+        heads = {}
+        for head in self.cfg.heads:
+            head_layers = []
+            in_feature = self.model.visual.output_dim
+
+            for i in range(head.n_hidden_layers):
+                head_layers += [
+                    nn.Linear(
+                        in_feature if i == 0 else self.cfg.hidden_features,
+                        self.cfg.hidden_features,
+                    ),
+                    self.make_activation(self.cfg.activation),
+                ]
+
+            head_layers = [nn.Sequential(*head_layers)]
+            head_layers += [
+                nn.Sequential(
+                    nn.Linear(
+                        self.cfg.hidden_features,
+                        self.cfg.hidden_features,
+                    ),
+                    self.make_activation(self.cfg.activation),
+                    nn.Linear(self.cfg.hidden_features, 1),
+                )
+                for _ in range(2)
+            ]
+            heads[head.name] = nn.ModuleList(head_layers)
+        self.heads = nn.ModuleDict(heads)
+
+    def make_activation(self, activation):
+        if activation == "relu":
+            return nn.ReLU(inplace=True)
+        elif activation == "silu":
+            return nn.SiLU(inplace=True)
+        else:
+            raise NotImplementedError
+
+    def forward(
+        self,
+        cond_image: Float[Tensor, "B 1 H W 4"],
+        sample: bool = True,
+    ) -> dict[str, Any]:
+        # Run the model
+        # Resize cond_image to 224
+        cond_image = cond_image.flatten(0, 1)
+        cond_image = nn.functional.interpolate(
+            cond_image.permute(0, 3, 1, 2),
+            size=(self.img_size, self.img_size),
+            mode="bilinear",
+            align_corners=False,
+        )
+        mask = cond_image[:, 3:4]
+        cond_image = cond_image[:, :3] * mask
+        cond_image = Normalize(
+            mean=OPENAI_DATASET_MEAN,
+            std=OPENAI_DATASET_STD,
+        )(cond_image)
+        mask = Normalize(0.5, 0.26)(mask).half()
+        image_features = self.model.visual(cond_image.half(), mask).float()
+
+        # Run the heads
+        outputs = {}
+
+        for head_dict in self.cfg.heads:
+            head_name = head_dict.name
+            shared_head, d1_h, d2_h = self.heads[head_name]
+            shared_features = shared_head(image_features)
+            d1, d2 = [head(shared_features).squeeze(-1) for head in [d1_h, d2_h]]
+            if self.cfg.distribution == "normal":
+                mean = d1
+                var = d2
+                if mean.shape[-1] == 1:
+                    outputs[head_name] = torch.distributions.Normal(
+                        mean + head_dict.output_bias,
+                        torch.nn.functional.softplus(var),
+                    )
+                else:
+                    outputs[head_name] = torch.distributions.MultivariateNormal(
+                        mean + head_dict.output_bias,
+                        torch.nn.functional.softplus(var).diag_embed(),
+                    )
+            elif self.cfg.distribution == "beta":
+                outputs[head_name] = torch.distributions.Beta(
+                    torch.nn.functional.softplus(d1 + head_dict.output_bias),
+                    torch.nn.functional.softplus(d2 + head_dict.output_bias),
+                )
+            else:
+                raise NotImplementedError
+
+        if sample:
+            for head_dict in self.cfg.heads:
+                head_name = head_dict.name
+                dist = outputs[head_name]
+
+                if head_dict.distribution_eval == "mean":
+                    out = dist.mean
+                elif head_dict.distribution_eval == "mode":
+                    out = dist.mode
+                elif head_dict.distribution_eval == "sample_mean":
+                    out = dist.sample([10]).mean(-1)
+                else:
+                    # use rsample if gradient is needed
+                    out = dist.rsample() if self.training else dist.sample()
+
+                outputs[head_name] = get_activation(head_dict.output_activation)(out)
+                outputs[f"{head_name}_dist"] = dist
+
+        for head in self.cfg.heads:
+            if head.shape:
+                if not sample:
+                    raise ValueError(
+                        "Cannot reshape non-sampled probabilisitic outputs"
+                    )
+                outputs[head.name] = outputs[head.name].reshape(*head.shape)
+
+            if head.add_to_decoder_features:
+                outputs[f"decoder_{head.name}"] = outputs[head.name]
+                del outputs[head.name]
+
+        return outputs

spar3d/models/isosurface.py

@@ -0,0 +1,229 @@
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from jaxtyping import Float, Integer
+from torch import Tensor
+
+from .mesh import Mesh
+
+
+class IsosurfaceHelper(nn.Module):
+    points_range: Tuple[float, float] = (0, 1)
+
+    @property
+    def grid_vertices(self) -> Float[Tensor, "N 3"]:
+        raise NotImplementedError
+
+    @property
+    def requires_instance_per_batch(self) -> bool:
+        return False
+
+
+class MarchingTetrahedraHelper(IsosurfaceHelper):
+    def __init__(self, resolution: int, tets_path: str):
+        super().__init__()
+        self.resolution = resolution
+        self.tets_path = tets_path
+
+        self.triangle_table: Float[Tensor, "..."]
+        self.register_buffer(
+            "triangle_table",
+            torch.as_tensor(
+                [
+                    [-1, -1, -1, -1, -1, -1],
+                    [1, 0, 2, -1, -1, -1],
+                    [4, 0, 3, -1, -1, -1],
+                    [1, 4, 2, 1, 3, 4],
+                    [3, 1, 5, -1, -1, -1],
+                    [2, 3, 0, 2, 5, 3],
+                    [1, 4, 0, 1, 5, 4],
+                    [4, 2, 5, -1, -1, -1],
+                    [4, 5, 2, -1, -1, -1],
+                    [4, 1, 0, 4, 5, 1],
+                    [3, 2, 0, 3, 5, 2],
+                    [1, 3, 5, -1, -1, -1],
+                    [4, 1, 2, 4, 3, 1],
+                    [3, 0, 4, -1, -1, -1],
+                    [2, 0, 1, -1, -1, -1],
+                    [-1, -1, -1, -1, -1, -1],
+                ],
+                dtype=torch.long,
+            ),
+            persistent=False,
+        )
+        self.num_triangles_table: Integer[Tensor, "..."]
+        self.register_buffer(
+            "num_triangles_table",
+            torch.as_tensor(
+                [0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long
+            ),
+            persistent=False,
+        )
+        self.base_tet_edges: Integer[Tensor, "..."]
+        self.register_buffer(
+            "base_tet_edges",
+            torch.as_tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long),
+            persistent=False,
+        )
+
+        tets = np.load(self.tets_path)
+        self._grid_vertices: Float[Tensor, "..."]
+        self.register_buffer(
+            "_grid_vertices",
+            torch.from_numpy(tets["vertices"]).float(),
+            persistent=False,
+        )
+        self.indices: Integer[Tensor, "..."]
+        self.register_buffer(
+            "indices", torch.from_numpy(tets["indices"]).long(), persistent=False
+        )
+
+        self._all_edges: Optional[Integer[Tensor, "Ne 2"]] = None
+
+        center_indices, boundary_indices = self.get_center_boundary_index(
+            self._grid_vertices
+        )
+        self.center_indices: Integer[Tensor, "..."]
+        self.register_buffer("center_indices", center_indices, persistent=False)
+        self.boundary_indices: Integer[Tensor, "..."]
+        self.register_buffer("boundary_indices", boundary_indices, persistent=False)
+
+    def get_center_boundary_index(self, verts):
+        magn = torch.sum(verts**2, dim=-1)
+
+        center_idx = torch.argmin(magn)
+        boundary_neg = verts == verts.max()
+        boundary_pos = verts == verts.min()
+
+        boundary = torch.bitwise_or(boundary_pos, boundary_neg)
+        boundary = torch.sum(boundary.float(), dim=-1)
+
+        boundary_idx = torch.nonzero(boundary)
+        return center_idx, boundary_idx.squeeze(dim=-1)
+
+    def normalize_grid_deformation(
+        self, grid_vertex_offsets: Float[Tensor, "Nv 3"]
+    ) -> Float[Tensor, "Nv 3"]:
+        return (
+            (self.points_range[1] - self.points_range[0])
+            / self.resolution  # half tet size is approximately 1 / self.resolution
+            * torch.tanh(grid_vertex_offsets)
+        )  # FIXME: hard-coded activation
+
+    @property
+    def grid_vertices(self) -> Float[Tensor, "Nv 3"]:
+        return self._grid_vertices
+
+    @property
+    def all_edges(self) -> Integer[Tensor, "Ne 2"]:
+        if self._all_edges is None:
+            # compute edges on GPU, or it would be VERY SLOW (basically due to the unique operation)
+            edges = torch.tensor(
+                [0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3],
+                dtype=torch.long,
+                device=self.indices.device,
+            )
+            _all_edges = self.indices[:, edges].reshape(-1, 2)
+            _all_edges_sorted = torch.sort(_all_edges, dim=1)[0]
+            _all_edges = torch.unique(_all_edges_sorted, dim=0)
+            self._all_edges = _all_edges
+        return self._all_edges
+
+    def sort_edges(self, edges_ex2):
+        with torch.no_grad():
+            order = (edges_ex2[:, 0] > edges_ex2[:, 1]).long()
+            order = order.unsqueeze(dim=1)
+
+            a = torch.gather(input=edges_ex2, index=order, dim=1)
+            b = torch.gather(input=edges_ex2, index=1 - order, dim=1)
+
+        return torch.stack([a, b], -1)
+
+    def _forward(self, pos_nx3, sdf_n, tet_fx4):
+        with torch.no_grad():
+            occ_n = sdf_n > 0
+            occ_fx4 = occ_n[tet_fx4.reshape(-1)].reshape(-1, 4)
+            occ_sum = torch.sum(occ_fx4, -1)
+            valid_tets = (occ_sum > 0) & (occ_sum < 4)
+            occ_sum = occ_sum[valid_tets]
+
+            # find all vertices
+            all_edges = tet_fx4[valid_tets][:, self.base_tet_edges].reshape(-1, 2)
+            all_edges = self.sort_edges(all_edges)
+            unique_edges, idx_map = torch.unique(all_edges, dim=0, return_inverse=True)
+
+            unique_edges = unique_edges.long()
+            mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1
+            mapping = (
+                torch.ones(
+                    (unique_edges.shape[0]), dtype=torch.long, device=pos_nx3.device
+                )
+                * -1
+            )
+            mapping[mask_edges] = torch.arange(
+                mask_edges.sum(), dtype=torch.long, device=pos_nx3.device
+            )
+            idx_map = mapping[idx_map]  # map edges to verts
+
+            interp_v = unique_edges[mask_edges]
+        edges_to_interp = pos_nx3[interp_v.reshape(-1)].reshape(-1, 2, 3)
+        edges_to_interp_sdf = sdf_n[interp_v.reshape(-1)].reshape(-1, 2, 1)
+        edges_to_interp_sdf[:, -1] *= -1
+
+        denominator = edges_to_interp_sdf.sum(1, keepdim=True)
+
+        edges_to_interp_sdf = torch.flip(edges_to_interp_sdf, [1]) / denominator
+        verts = (edges_to_interp * edges_to_interp_sdf).sum(1)
+
+        idx_map = idx_map.reshape(-1, 6)
+
+        v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=pos_nx3.device))
+        tetindex = (occ_fx4[valid_tets] * v_id.unsqueeze(0)).sum(-1)
+        num_triangles = self.num_triangles_table[tetindex]
+
+        # Generate triangle indices
+        faces = torch.cat(
+            (
+                torch.gather(
+                    input=idx_map[num_triangles == 1],
+                    dim=1,
+                    index=self.triangle_table[tetindex[num_triangles == 1]][:, :3],
+                ).reshape(-1, 3),
+                torch.gather(
+                    input=idx_map[num_triangles == 2],
+                    dim=1,
+                    index=self.triangle_table[tetindex[num_triangles == 2]][:, :6],
+                ).reshape(-1, 3),
+            ),
+            dim=0,
+        )
+
+        return verts, faces
+
+    def forward(
+        self,
+        level: Float[Tensor, "N3 1"],
+        deformation: Optional[Float[Tensor, "N3 3"]] = None,
+    ) -> Mesh:
+        if deformation is not None:
+            grid_vertices = self.grid_vertices + self.normalize_grid_deformation(
+                deformation
+            )
+        else:
+            grid_vertices = self.grid_vertices
+
+        v_pos, t_pos_idx = self._forward(grid_vertices, level, self.indices)
+
+        mesh = Mesh(
+            v_pos=v_pos,
+            t_pos_idx=t_pos_idx,
+            # extras
+            grid_vertices=grid_vertices,
+            tet_edges=self.all_edges,
+            grid_level=level,
+            grid_deformation=deformation,
+        )
+
+        return mesh

spar3d/models/mesh.py

@@ -0,0 +1,317 @@
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import trimesh
+from jaxtyping import Float, Integer
+from torch import Tensor
+
+from spar3d.models.utils import dot
+
+try:
+    from uv_unwrapper import Unwrapper
+except ImportError:
+    import logging
+
+    logging.warning(
+        "Could not import uv_unwrapper. Please install it via `pip install uv_unwrapper/`"
+    )
+    # Exit early to avoid further errors
+    raise ImportError("uv_unwrapper not found")
+
+try:
+    import gpytoolbox
+
+    TRIANGLE_REMESH_AVAILABLE = True
+except ImportError:
+    TRIANGLE_REMESH_AVAILABLE = False
+    import logging
+
+    logging.warning(
+        "Could not import gpytoolbox. Triangle remeshing functionality will be disabled. "
+        "Install via `pip install gpytoolbox`"
+    )
+
+try:
+    import pynim
+
+    QUAD_REMESH_AVAILABLE = True
+except ImportError:
+    QUAD_REMESH_AVAILABLE = False
+    import logging
+
+    logging.warning(
+        "Could not import pynim. Quad remeshing functionality will be disabled. "
+        "Install via `pip install git+https://github.com/vork/[email protected]`"
+    )
+
+
+class Mesh:
+    def __init__(
+        self, v_pos: Float[Tensor, "Nv 3"], t_pos_idx: Integer[Tensor, "Nf 3"], **kwargs
+    ) -> None:
+        self.v_pos: Float[Tensor, "Nv 3"] = v_pos
+        self.t_pos_idx: Integer[Tensor, "Nf 3"] = t_pos_idx
+        self._v_nrm: Optional[Float[Tensor, "Nv 3"]] = None
+        self._v_tng: Optional[Float[Tensor, "Nv 3"]] = None
+        self._v_tex: Optional[Float[Tensor, "Nt 3"]] = None
+        self._edges: Optional[Integer[Tensor, "Ne 2"]] = None
+        self.extras: Dict[str, Any] = {}
+        for k, v in kwargs.items():
+            self.add_extra(k, v)
+
+        self.unwrapper = Unwrapper()
+
+    def add_extra(self, k, v) -> None:
+        self.extras[k] = v
+
+    @property
+    def requires_grad(self):
+        return self.v_pos.requires_grad
+
+    @property
+    def v_nrm(self):
+        if self._v_nrm is None:
+            self._v_nrm = self._compute_vertex_normal()
+        return self._v_nrm
+
+    @property
+    def v_tng(self):
+        if self._v_tng is None:
+            self._v_tng = self._compute_vertex_tangent()
+        return self._v_tng
+
+    @property
+    def v_tex(self):
+        if self._v_tex is None:
+            self.unwrap_uv()
+        return self._v_tex
+
+    @property
+    def edges(self):
+        if self._edges is None:
+            self._edges = self._compute_edges()
+        return self._edges
+
+    def _compute_vertex_normal(self):
+        i0 = self.t_pos_idx[:, 0]
+        i1 = self.t_pos_idx[:, 1]
+        i2 = self.t_pos_idx[:, 2]
+
+        v0 = self.v_pos[i0, :]
+        v1 = self.v_pos[i1, :]
+        v2 = self.v_pos[i2, :]
+
+        face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
+
+        # Splat face normals to vertices
+        v_nrm = torch.zeros_like(self.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(
+            dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.as_tensor([0.0, 0.0, 1.0]).to(v_nrm)
+        )
+        v_nrm = F.normalize(v_nrm, dim=1)
+
+        if torch.is_anomaly_enabled():
+            assert torch.all(torch.isfinite(v_nrm))
+
+        return v_nrm
+
+    def _compute_vertex_tangent(self):
+        vn_idx = [None] * 3
+        pos = [None] * 3
+        tex = [None] * 3
+        for i in range(0, 3):
+            pos[i] = self.v_pos[self.t_pos_idx[:, i]]
+            tex[i] = self.v_tex[self.t_pos_idx[:, i]]
+            # t_nrm_idx is always the same as t_pos_idx
+            vn_idx[i] = self.t_pos_idx[:, i]
+
+        tangents = torch.zeros_like(self.v_nrm)
+        tansum = torch.zeros_like(self.v_nrm)
+
+        # Compute tangent space for each triangle
+        duv1 = tex[1] - tex[0]
+        duv2 = tex[2] - tex[0]
+        dpos1 = pos[1] - pos[0]
+        dpos2 = pos[2] - pos[0]
+
+        tng_nom = dpos1 * duv2[..., 1:2] - dpos2 * duv1[..., 1:2]
+
+        denom = duv1[..., 0:1] * duv2[..., 1:2] - duv1[..., 1:2] * duv2[..., 0:1]
+
+        # Avoid division by zero for degenerated texture coordinates
+        denom_safe = denom.clip(1e-6)
+        tang = tng_nom / denom_safe
+
+        # 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
+        # Also normalize it. Here we do not normalize the individual triangles first so larger area
+        # triangles influence the tangent space more
+        tangents = tangents / tansum
+
+        # Normalize and make sure tangent is perpendicular to normal
+        tangents = F.normalize(tangents, dim=1)
+        tangents = F.normalize(tangents - dot(tangents, self.v_nrm) * self.v_nrm)
+
+        if torch.is_anomaly_enabled():
+            assert torch.all(torch.isfinite(tangents))
+
+        return tangents
+
+    def quad_remesh(
+        self,
+        quad_vertex_count: int = -1,
+        quad_rosy: int = 4,
+        quad_crease_angle: float = -1.0,
+        quad_smooth_iter: int = 2,
+        quad_align_to_boundaries: bool = False,
+    ) -> Mesh:
+        if not QUAD_REMESH_AVAILABLE:
+            raise ImportError("Quad remeshing requires pynim to be installed")
+        if quad_vertex_count < 0:
+            quad_vertex_count = self.v_pos.shape[0]
+        v_pos = self.v_pos.detach().cpu().numpy().astype(np.float32)
+        t_pos_idx = self.t_pos_idx.detach().cpu().numpy().astype(np.uint32)
+
+        new_vert, new_faces = pynim.remesh(
+            v_pos,
+            t_pos_idx,
+            quad_vertex_count // 4,
+            rosy=quad_rosy,
+            posy=4,
+            creaseAngle=quad_crease_angle,
+            align_to_boundaries=quad_align_to_boundaries,
+            smooth_iter=quad_smooth_iter,
+            deterministic=False,
+        )
+
+        # Briefly load in trimesh
+        mesh = trimesh.Trimesh(vertices=new_vert, faces=new_faces.astype(np.int32))
+
+        v_pos = torch.from_numpy(mesh.vertices).to(self.v_pos).contiguous()
+        t_pos_idx = torch.from_numpy(mesh.faces).to(self.t_pos_idx).contiguous()
+
+        # Create new mesh
+        return Mesh(v_pos, t_pos_idx)
+
+    def triangle_remesh(
+        self,
+        triangle_average_edge_length_multiplier: Optional[float] = None,
+        triangle_remesh_steps: int = 10,
+        triangle_vertex_count=-1,
+    ):
+        if not TRIANGLE_REMESH_AVAILABLE:
+            raise ImportError("Triangle remeshing requires gpytoolbox to be installed")
+        if triangle_vertex_count > 0:
+            reduction = triangle_vertex_count / self.v_pos.shape[0]
+            print("Triangle reduction:", reduction)
+            v_pos = self.v_pos.detach().cpu().numpy().astype(np.float32)
+            t_pos_idx = self.t_pos_idx.detach().cpu().numpy().astype(np.int32)
+            if reduction > 1.0:
+                subdivide_iters = int(math.ceil(math.log(reduction) / math.log(2)))
+                print("Subdivide iters:", subdivide_iters)
+                v_pos, t_pos_idx = gpytoolbox.subdivide(
+                    v_pos,
+                    t_pos_idx,
+                    iters=subdivide_iters,
+                )
+                reduction = triangle_vertex_count / v_pos.shape[0]
+
+            # Simplify
+            points_out, faces_out, _, _ = gpytoolbox.decimate(
+                v_pos,
+                t_pos_idx,
+                face_ratio=reduction,
+            )
+
+            # Convert back to torch
+            self.v_pos = torch.from_numpy(points_out).to(self.v_pos)
+            self.t_pos_idx = torch.from_numpy(faces_out).to(self.t_pos_idx)
+            self._edges = None
+            triangle_average_edge_length_multiplier = None
+
+        edges = self.edges
+        if triangle_average_edge_length_multiplier is None:
+            h = None
+        else:
+            h = float(
+                torch.linalg.norm(
+                    self.v_pos[edges[:, 0]] - self.v_pos[edges[:, 1]], dim=1
+                )
+                .mean()
+                .item()
+                * triangle_average_edge_length_multiplier
+            )
+
+        # Convert to numpy
+        v_pos = self.v_pos.detach().cpu().numpy().astype(np.float64)
+        t_pos_idx = self.t_pos_idx.detach().cpu().numpy().astype(np.int32)
+
+        # Remesh
+        v_remesh, f_remesh = gpytoolbox.remesh_botsch(
+            v_pos,
+            t_pos_idx,
+            triangle_remesh_steps,
+            h,
+        )
+
+        # Convert back to torch
+        v_pos = torch.from_numpy(v_remesh).to(self.v_pos).contiguous()
+        t_pos_idx = torch.from_numpy(f_remesh).to(self.t_pos_idx).contiguous()
+
+        # Create new mesh
+        return Mesh(v_pos, t_pos_idx)
+
+    @torch.no_grad()
+    def unwrap_uv(
+        self,
+        island_padding: float = 0.02,
+    ) -> Mesh:
+        uv, indices = self.unwrapper(
+            self.v_pos, self.v_nrm, self.t_pos_idx, island_padding
+        )
+
+        # Do store per vertex UVs.
+        # This means we need to duplicate some vertices at the seams
+        individual_vertices = self.v_pos[self.t_pos_idx].reshape(-1, 3)
+        individual_faces = torch.arange(
+            individual_vertices.shape[0],
+            device=individual_vertices.device,
+            dtype=self.t_pos_idx.dtype,
+        ).reshape(-1, 3)
+        uv_flat = uv[indices].reshape((-1, 2))
+        # uv_flat[:, 1] = 1 - uv_flat[:, 1]
+
+        self.v_pos = individual_vertices
+        self.t_pos_idx = individual_faces
+        self._v_tex = uv_flat
+        self._v_nrm = self._compute_vertex_normal()
+        self._v_tng = self._compute_vertex_tangent()
+
+    def _compute_edges(self):
+        # Compute edges
+        edges = torch.cat(
+            [
+                self.t_pos_idx[:, [0, 1]],
+                self.t_pos_idx[:, [1, 2]],
+                self.t_pos_idx[:, [2, 0]],
+            ],
+            dim=0,
+        )
+        edges = edges.sort()[0]
+        edges = torch.unique(edges, dim=0)
+        return edges

spar3d/models/network.py

@@ -0,0 +1,223 @@
+from dataclasses import dataclass, field
+from typing import Callable, List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from jaxtyping import Float
+from torch import Tensor
+from torch.autograd import Function
+from torch.cuda.amp import custom_bwd, custom_fwd
+
+from spar3d.models.utils import BaseModule, normalize
+from spar3d.utils import get_device
+
+
+def conditional_decorator(decorator_with_args, condition, *args, **kwargs):
+    def wrapper(fn):
+        if condition:
+            if len(kwargs) == 0:
+                return decorator_with_args
+            return decorator_with_args(*args, **kwargs)(fn)
+        else:
+            return fn
+
+    return wrapper
+
+
+class PixelShuffleUpsampleNetwork(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int = 1024
+        out_channels: int = 40
+        scale_factor: int = 4
+
+        conv_layers: int = 4
+        conv_kernel_size: int = 3
+
+    cfg: Config
+
+    def configure(self) -> None:
+        layers = []
+        output_channels = self.cfg.out_channels * self.cfg.scale_factor**2
+
+        in_channels = self.cfg.in_channels
+        for i in range(self.cfg.conv_layers):
+            cur_out_channels = (
+                in_channels if i != self.cfg.conv_layers - 1 else output_channels
+            )
+            layers.append(
+                nn.Conv2d(
+                    in_channels,
+                    cur_out_channels,
+                    self.cfg.conv_kernel_size,
+                    padding=(self.cfg.conv_kernel_size - 1) // 2,
+                )
+            )
+            if i != self.cfg.conv_layers - 1:
+                layers.append(nn.ReLU(inplace=True))
+
+        layers.append(nn.PixelShuffle(self.cfg.scale_factor))
+
+        self.upsample = nn.Sequential(*layers)
+
+    def forward(
+        self, triplanes: Float[Tensor, "B 3 Ci Hp Wp"]
+    ) -> Float[Tensor, "B 3 Co Hp2 Wp2"]:
+        return rearrange(
+            self.upsample(
+                rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3)
+            ),
+            "(B Np) Co Hp Wp -> B Np Co Hp Wp",
+            Np=3,
+        )
+
+
+class _TruncExp(Function):  # pylint: disable=abstract-method
+    # Implementation from torch-ngp:
+    # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
+    @staticmethod
+    @conditional_decorator(
+        custom_fwd, "cuda" in get_device(), cast_inputs=torch.float32
+    )
+    def forward(ctx, x):  # pylint: disable=arguments-differ
+        ctx.save_for_backward(x)
+        return torch.exp(x)
+
+    @staticmethod
+    @conditional_decorator(custom_bwd, "cuda" in get_device())
+    def backward(ctx, g):  # pylint: disable=arguments-differ
+        x = ctx.saved_tensors[0]
+        return g * torch.exp(torch.clamp(x, max=15))
+
+
+trunc_exp = _TruncExp.apply
+
+
+def get_activation(name) -> Callable:
+    if name is None:
+        return lambda x: x
+    name = name.lower()
+    if name == "none" or name == "linear" or name == "identity":
+        return lambda x: x
+    elif name == "lin2srgb":
+        return lambda x: torch.where(
+            x > 0.0031308,
+            torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) * 1.055 - 0.055,
+            12.92 * x,
+        ).clamp(0.0, 1.0)
+    elif name == "exp":
+        return lambda x: torch.exp(x)
+    elif name == "shifted_exp":
+        return lambda x: torch.exp(x - 1.0)
+    elif name == "trunc_exp":
+        return trunc_exp
+    elif name == "shifted_trunc_exp":
+        return lambda x: trunc_exp(x - 1.0)
+    elif name == "sigmoid":
+        return lambda x: torch.sigmoid(x)
+    elif name == "tanh":
+        return lambda x: torch.tanh(x)
+    elif name == "shifted_softplus":
+        return lambda x: F.softplus(x - 1.0)
+    elif name == "scale_-11_01":
+        return lambda x: x * 0.5 + 0.5
+    elif name == "negative":
+        return lambda x: -x
+    elif name == "normalize_channel_last":
+        return lambda x: normalize(x)
+    elif name == "normalize_channel_first":
+        return lambda x: normalize(x, dim=1)
+    else:
+        try:
+            return getattr(F, name)
+        except AttributeError:
+            raise ValueError(f"Unknown activation function: {name}")
+
+
+class LambdaModule(torch.nn.Module):
+    def __init__(self, lambd: Callable[[torch.Tensor], torch.Tensor]):
+        super().__init__()
+        self.lambd = lambd
+
+    def forward(self, x):
+        return self.lambd(x)
+
+
+def get_activation_module(name) -> torch.nn.Module:
+    return LambdaModule(get_activation(name))
+
+
+@dataclass
+class HeadSpec:
+    name: str
+    out_channels: int
+    n_hidden_layers: int
+    output_activation: Optional[str] = None
+    out_bias: float = 0.0
+
+
+class MaterialMLP(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int = 120
+        n_neurons: int = 64
+        activation: str = "silu"
+        heads: List[HeadSpec] = field(default_factory=lambda: [])
+
+    cfg: Config
+
+    def configure(self) -> None:
+        assert len(self.cfg.heads) > 0
+        heads = {}
+        for head in self.cfg.heads:
+            head_layers = []
+            for i in range(head.n_hidden_layers):
+                head_layers += [
+                    nn.Linear(
+                        self.cfg.in_channels if i == 0 else self.cfg.n_neurons,
+                        self.cfg.n_neurons,
+                    ),
+                    self.make_activation(self.cfg.activation),
+                ]
+            head_layers += [
+                nn.Linear(
+                    self.cfg.n_neurons,
+                    head.out_channels,
+                ),
+            ]
+            heads[head.name] = nn.Sequential(*head_layers)
+        self.heads = nn.ModuleDict(heads)
+
+    def make_activation(self, activation):
+        if activation == "relu":
+            return nn.ReLU(inplace=True)
+        elif activation == "silu":
+            return nn.SiLU(inplace=True)
+        else:
+            raise NotImplementedError
+
+    def keys(self):
+        return self.heads.keys()
+
+    def forward(
+        self, x, include: Optional[List] = None, exclude: Optional[List] = None
+    ):
+        if include is not None and exclude is not None:
+            raise ValueError("Cannot specify both include and exclude.")
+        if include is not None:
+            heads = [h for h in self.cfg.heads if h.name in include]
+        elif exclude is not None:
+            heads = [h for h in self.cfg.heads if h.name not in exclude]
+        else:
+            heads = self.cfg.heads
+
+        out = {
+            head.name: get_activation(head.output_activation)(
+                self.heads[head.name](x) + head.out_bias
+            )
+            for head in heads
+        }
+
+        return out

spar3d/models/tokenizers/dinov2.py

@@ -0,0 +1,1196 @@
+# coding=utf-8
+# Copyright 2023 Meta AI and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch DINOv2 model."""
+
+import collections.abc
+import math
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Set, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (
+    BackboneOutput,
+    BaseModelOutput,
+    BaseModelOutputWithPooling,
+    ImageClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.dinov2.configuration_dinov2 import Dinov2Config
+from transformers.pytorch_utils import (
+    find_pruneable_heads_and_indices,
+    prune_linear_layer,
+)
+from transformers.utils import (
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.utils.backbone_utils import BackboneMixin
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "Dinov2Config"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "facebook/dinov2-base"
+_EXPECTED_OUTPUT_SHAPE = [1, 257, 768]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "facebook/dinov2-base"
+
+
+DINOV2_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "facebook/dinov2-base",
+    # See all DINOv2 models at https://huggingface.co/models?filter=dinov2
+]
+
+
+class Dinov2Embeddings(nn.Module):
+    """
+    Construct the CLS token, mask token, position and patch embeddings.
+    """
+
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+
+        self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+        # register as mask token as it's not used in optimization
+        # to avoid the use of find_unused_parameters_true
+        # self.mask_token = nn.Parameter(torch.zeros(1, config.hidden_size))
+        self.register_buffer("mask_token", torch.zeros(1, config.hidden_size))
+        self.patch_embeddings = Dinov2PatchEmbeddings(config)
+        num_patches = self.patch_embeddings.num_patches
+        self.position_embeddings = nn.Parameter(
+            torch.randn(1, num_patches + 1, config.hidden_size)
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.config = config
+
+    def interpolate_pos_encoding(
+        self, embeddings: torch.Tensor, height: int, width: int
+    ) -> torch.Tensor:
+        """
+        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
+        resolution images.
+
+        Source:
+        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
+        """
+
+        num_patches = embeddings.shape[1] - 1
+        num_positions = self.position_embeddings.shape[1] - 1
+        if num_patches == num_positions and height == width:
+            return self.position_embeddings
+        class_pos_embed = self.position_embeddings[:, 0]
+        patch_pos_embed = self.position_embeddings[:, 1:]
+        dim = embeddings.shape[-1]
+        height = height // self.config.patch_size
+        width = width // self.config.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        height, width = height + 0.1, width + 0.1
+        patch_pos_embed = patch_pos_embed.reshape(
+            1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim
+        )
+        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed,
+            scale_factor=(
+                height / math.sqrt(num_positions),
+                width / math.sqrt(num_positions),
+            ),
+            mode="bicubic",
+            align_corners=False,
+        )
+        if (
+            int(height) != patch_pos_embed.shape[-2]
+            or int(width) != patch_pos_embed.shape[-1]
+        ):
+            raise ValueError(
+                "Width or height does not match with the interpolated position embeddings"
+            )
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        bool_masked_pos: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        batch_size, _, height, width = pixel_values.shape
+        patch_embeddings = self.patch_embeddings(pixel_values)
+        embeddings = patch_embeddings
+
+        if bool_masked_pos is not None:
+            embeddings = torch.where(
+                bool_masked_pos.unsqueeze(-1),
+                self.mask_token.to(embeddings.dtype).unsqueeze(0),
+                embeddings,
+            )
+
+        # add the [CLS] token to the embedded patch tokens
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        embeddings = torch.cat((cls_tokens, embeddings), dim=1)
+
+        # add positional encoding to each token
+        embeddings = embeddings + self.interpolate_pos_encoding(
+            embeddings, height, width
+        )
+
+        embeddings = self.dropout(embeddings)
+
+        return embeddings
+
+
+class Dinov2PatchEmbeddings(nn.Module):
+    """
+    This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+    `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
+    Transformer.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        image_size, patch_size = config.image_size, config.patch_size
+        num_channels, hidden_size = config.num_channels, config.hidden_size
+
+        image_size = (
+            image_size
+            if isinstance(image_size, collections.abc.Iterable)
+            else (image_size, image_size)
+        )
+        patch_size = (
+            patch_size
+            if isinstance(patch_size, collections.abc.Iterable)
+            else (patch_size, patch_size)
+        )
+        num_patches = (image_size[1] // patch_size[1]) * (
+            image_size[0] // patch_size[0]
+        )
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.num_patches = num_patches
+
+        self.projection = nn.Conv2d(
+            num_channels, hidden_size, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        num_channels = pixel_values.shape[1]
+        if num_channels != self.num_channels:
+            raise ValueError(
+                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+                f" Expected {self.num_channels} but got {num_channels}."
+            )
+        """
+        embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
+        return embeddings
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->Dinov2
+class Dinov2SelfAttention(nn.Module):
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
+            config, "embedding_size"
+        ):
+            raise ValueError(
+                f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
+                f"heads {config.num_attention_heads}."
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.attention_probs_dropout_prob = config.attention_probs_dropout_prob
+
+        self.query = nn.Linear(
+            config.hidden_size, self.all_head_size, bias=config.qkv_bias
+        )
+        self.key = nn.Linear(
+            config.hidden_size, self.all_head_size, bias=config.qkv_bias
+        )
+        self.value = nn.Linear(
+            config.hidden_size, self.all_head_size, bias=config.qkv_bias
+        )
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (
+            self.num_attention_heads,
+            self.attention_head_size,
+        )
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        mixed_query_layer = self.query(hidden_states)
+
+        if hasattr(F, "scaled_dot_product_attention"):
+            assert head_mask is None and not output_attentions
+            new_size = hidden_states.size()[:-1] + (
+                self.num_attention_heads,
+                self.attention_head_size,
+            )
+            key_layer = self.key(hidden_states).reshape(new_size).transpose(1, 2)
+            value_layer = self.value(hidden_states).reshape(new_size).transpose(1, 2)
+            query_layer = mixed_query_layer.reshape(new_size).transpose(1, 2)
+            context_layer = F.scaled_dot_product_attention(
+                query_layer,
+                key_layer,
+                value_layer,
+                dropout_p=self.attention_probs_dropout_prob,
+                is_causal=False,
+            )
+            context_layer = context_layer.transpose(1, 2).reshape(
+                *hidden_states.size()[:-1], -1
+            )
+        else:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+            query_layer = self.transpose_for_scores(mixed_query_layer)
+
+            # Take the dot product between "query" and "key" to get the raw attention scores.
+            attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+            attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+
+            # Normalize the attention scores to probabilities.
+            attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+            # This is actually dropping out entire tokens to attend to, which might
+            # seem a bit unusual, but is taken from the original Transformer paper.
+            attention_probs = self.dropout(attention_probs)
+
+            # Mask heads if we want to
+            if head_mask is not None:
+                attention_probs = attention_probs * head_mask
+
+            context_layer = torch.matmul(attention_probs, value_layer)
+
+            context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+            new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+            context_layer = context_layer.view(new_context_layer_shape)
+
+        outputs = (
+            (context_layer, attention_probs) if output_attentions else (context_layer,)
+        )
+
+        return outputs
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->Dinov2
+class Dinov2SelfOutput(nn.Module):
+    """
+    The residual connection is defined in Dinov2Layer instead of here (as is the case with other models), due to the
+    layernorm applied before each block.
+    """
+
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(
+        self, hidden_states: torch.Tensor, input_tensor: torch.Tensor
+    ) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        return hidden_states
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->Dinov2
+class Dinov2Attention(nn.Module):
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+        self.attention = Dinov2SelfAttention(config)
+        self.output = Dinov2SelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads: Set[int]) -> None:
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads,
+            self.attention.num_attention_heads,
+            self.attention.attention_head_size,
+            self.pruned_heads,
+        )
+
+        # Prune linear layers
+        self.attention.query = prune_linear_layer(self.attention.query, index)
+        self.attention.key = prune_linear_layer(self.attention.key, index)
+        self.attention.value = prune_linear_layer(self.attention.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.attention.num_attention_heads = self.attention.num_attention_heads - len(
+            heads
+        )
+        self.attention.all_head_size = (
+            self.attention.attention_head_size * self.attention.num_attention_heads
+        )
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        self_outputs = self.attention(hidden_states, head_mask, output_attentions)
+
+        attention_output = self.output(self_outputs[0], hidden_states)
+
+        outputs = (attention_output,) + self_outputs[
+            1:
+        ]  # add attentions if we output them
+        return outputs
+
+
+class Dinov2LayerScale(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.lambda1 = nn.Parameter(
+            config.layerscale_value * torch.ones(config.hidden_size)
+        )
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        return hidden_state * self.lambda1
+
+
+# Copied from transformers.models.beit.modeling_beit.drop_path
+def drop_path(
+    input: torch.Tensor, drop_prob: float = 0.0, training: bool = False
+) -> torch.Tensor:
+    """
+    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
+    however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
+    layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
+    argument.
+    """
+    if drop_prob == 0.0 or not training:
+        return input
+    keep_prob = 1 - drop_prob
+    shape = (input.shape[0],) + (1,) * (
+        input.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(
+        shape, dtype=input.dtype, device=input.device
+    )
+    random_tensor.floor_()  # binarize
+    output = input.div(keep_prob) * random_tensor
+    return output
+
+
+# Copied from transformers.models.beit.modeling_beit.BeitDropPath
+class Dinov2DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob: Optional[float] = None) -> None:
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return drop_path(hidden_states, self.drop_prob, self.training)
+
+    def extra_repr(self) -> str:
+        return "p={}".format(self.drop_prob)
+
+
+class Dinov2MLP(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        in_features = out_features = config.hidden_size
+        hidden_features = int(config.hidden_size * config.mlp_ratio)
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=True)
+        if isinstance(config.hidden_act, str):
+            self.activation = ACT2FN[config.hidden_act]
+        else:
+            self.activation = config.hidden_act
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=True)
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        hidden_state = self.fc1(hidden_state)
+        hidden_state = self.activation(hidden_state)
+        hidden_state = self.fc2(hidden_state)
+        return hidden_state
+
+
+class Dinov2SwiGLUFFN(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        in_features = out_features = config.hidden_size
+        hidden_features = int(config.hidden_size * config.mlp_ratio)
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+
+        self.weights_in = nn.Linear(in_features, 2 * hidden_features, bias=True)
+        self.weights_out = nn.Linear(hidden_features, out_features, bias=True)
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        hidden_state = self.weights_in(hidden_state)
+        x1, x2 = hidden_state.chunk(2, dim=-1)
+        hidden = nn.functional.silu(x1) * x2
+        return self.weights_out(hidden)
+
+
+class Dinov2Layer(nn.Module):
+    """This corresponds to the Block class in the original implementation."""
+
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+
+        self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.norm1_modulation = None
+        self.attention = Dinov2Attention(config)
+        self.layer_scale1 = Dinov2LayerScale(config)
+        self.drop_path1 = (
+            Dinov2DropPath(config.drop_path_rate)
+            if config.drop_path_rate > 0.0
+            else nn.Identity()
+        )
+
+        self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.norm2_modulation = None
+
+        if config.use_swiglu_ffn:
+            self.mlp = Dinov2SwiGLUFFN(config)
+        else:
+            self.mlp = Dinov2MLP(config)
+        self.layer_scale2 = Dinov2LayerScale(config)
+        self.drop_path2 = (
+            Dinov2DropPath(config.drop_path_rate)
+            if config.drop_path_rate > 0.0
+            else nn.Identity()
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        head_mask: Optional[torch.Tensor] = None,
+        modulation_cond: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        hidden_states_norm = self.norm1(hidden_states)
+        if self.norm1_modulation is not None:
+            assert modulation_cond is not None
+            hidden_states_norm = self.norm1_modulation(
+                hidden_states_norm, modulation_cond
+            )
+        self_attention_outputs = self.attention(
+            hidden_states_norm,  # in Dinov2, layernorm is applied before self-attention
+            head_mask,
+            output_attentions=output_attentions,
+        )
+        attention_output = self_attention_outputs[0]
+
+        attention_output = self.layer_scale1(attention_output)
+        outputs = self_attention_outputs[
+            1:
+        ]  # add self attentions if we output attention weights
+
+        # first residual connection
+        hidden_states = attention_output + hidden_states
+
+        # in Dinov2, layernorm is also applied after self-attention
+        layer_output = self.norm2(hidden_states)
+        if self.norm2_modulation is not None:
+            assert modulation_cond is not None
+            layer_output = self.norm2_modulation(layer_output, modulation_cond)
+        layer_output = self.mlp(layer_output)
+        layer_output = self.layer_scale2(layer_output)
+
+        # second residual connection
+        layer_output = layer_output + hidden_states
+
+        outputs = (layer_output,) + outputs
+
+        return outputs
+
+    def register_ada_norm_modulation(self, norm1_mod: nn.Module, norm2_mod: nn.Module):
+        self.norm1_modulation = norm1_mod
+        self.norm2_modulation = norm2_mod
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTEncoder with ViT->Dinov2
+class Dinov2Encoder(nn.Module):
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList(
+            [Dinov2Layer(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        head_mask: Optional[torch.Tensor] = None,
+        modulation_cond: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ) -> Union[tuple, BaseModelOutput]:
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+
+        for i, layer_module in enumerate(self.layer):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    layer_head_mask,
+                    modulation_cond,
+                    use_reentrant=False,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states, layer_head_mask, modulation_cond, output_attentions
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, all_hidden_states, all_self_attentions]
+                if v is not None
+            )
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+
+class Dinov2PreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = Dinov2Config
+    base_model_prefix = "dinov2"
+    main_input_name = "pixel_values"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None:
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid
+            # `trunc_normal_cpu` not implemented in `half` issues
+            module.weight.data = nn.init.trunc_normal_(
+                module.weight.data.to(torch.float32),
+                mean=0.0,
+                std=self.config.initializer_range,
+            ).to(module.weight.dtype)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, Dinov2Embeddings):
+            module.position_embeddings.data = nn.init.trunc_normal_(
+                module.position_embeddings.data.to(torch.float32),
+                mean=0.0,
+                std=self.config.initializer_range,
+            ).to(module.position_embeddings.dtype)
+
+            module.cls_token.data = nn.init.trunc_normal_(
+                module.cls_token.data.to(torch.float32),
+                mean=0.0,
+                std=self.config.initializer_range,
+            ).to(module.cls_token.dtype)
+
+    def _set_gradient_checkpointing(
+        self, module: Dinov2Encoder, value: bool = False
+    ) -> None:
+        if isinstance(module, Dinov2Encoder):
+            module.gradient_checkpointing = value
+
+
+DINOV2_START_DOCSTRING = r"""
+    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
+    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
+    behavior.
+
+    Parameters:
+        config ([`Dinov2Config`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+DINOV2_BASE_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
+            [`BitImageProcessor.preprocess`] for details.
+
+        bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`):
+            Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for
+            pre-training.
+
+        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
+            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+DINOV2_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
+            [`BitImageProcessor.preprocess`] for details.
+
+        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
+            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@dataclass
+class CustomBaseModelOutputWithPooling(BaseModelOutputWithPooling):
+    patch_embeddings: Optional[torch.FloatTensor] = None
+
+
+@add_start_docstrings(
+    "The bare DINOv2 Model transformer outputting raw hidden-states without any specific head on top.",
+    DINOV2_START_DOCSTRING,
+)
+class Dinov2Model(Dinov2PreTrainedModel):
+    def __init__(self, config: Dinov2Config):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = Dinov2Embeddings(config)
+        self.encoder = Dinov2Encoder(config)
+
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> Dinov2PatchEmbeddings:
+        return self.embeddings.patch_embeddings
+
+    def expand_input_channels(self, extra_input_channels: int) -> None:
+        if extra_input_channels == 0:
+            return
+        conv_old = self.embeddings.patch_embeddings.projection
+        conv_new = nn.Conv2d(
+            self.config.num_channels + extra_input_channels,
+            self.config.hidden_size,
+            kernel_size=self.config.patch_size,
+            stride=self.config.patch_size,
+        ).to(self.device)
+        with torch.no_grad():
+            conv_new.weight[:, :3] = conv_old.weight
+            conv_new.bias = conv_old.bias
+        self.embeddings.patch_embeddings.projection = conv_new
+        del conv_old
+
+    def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None:
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    @add_start_docstrings_to_model_forward(DINOV2_BASE_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=BaseModelOutputWithPooling,
+        config_class=_CONFIG_FOR_DOC,
+        modality="vision",
+        expected_output=_EXPECTED_OUTPUT_SHAPE,
+    )
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        bool_masked_pos: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        modulation_cond: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if pixel_values is None:
+            raise ValueError("You have to specify pixel_values")
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        embedding_output = self.embeddings(
+            pixel_values, bool_masked_pos=bool_masked_pos
+        )
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            head_mask=head_mask,
+            modulation_cond=modulation_cond,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        sequence_output = self.layernorm(sequence_output)
+        pooled_output = sequence_output[:, 0, :]
+
+        if not return_dict:
+            head_outputs = (sequence_output, pooled_output)
+            return head_outputs + encoder_outputs[1:]
+
+        return CustomBaseModelOutputWithPooling(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            patch_embeddings=embedding_output,
+        )
+
+    def set_gradient_checkpointing(self, value: bool = False) -> None:
+        self._set_gradient_checkpointing(self.encoder, value)
+
+
+@add_start_docstrings(
+    """
+    Dinov2 Model transformer with an image classification head on top (a linear layer on top of the final hidden state
+    of the [CLS] token) e.g. for ImageNet.
+    """,
+    DINOV2_START_DOCSTRING,
+)
+class Dinov2ForImageClassification(Dinov2PreTrainedModel):
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__(config)
+
+        self.num_labels = config.num_labels
+        self.dinov2 = Dinov2Model(config)
+
+        # Classifier head
+        self.classifier = (
+            nn.Linear(config.hidden_size * 2, config.num_labels)
+            if config.num_labels > 0
+            else nn.Identity()
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_IMAGE_CLASS_CHECKPOINT,
+        output_type=ImageClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[tuple, ImageClassifierOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = self.dinov2(
+            pixel_values,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]  # batch_size, sequence_length, hidden_size
+
+        cls_token = sequence_output[:, 0]
+        patch_tokens = sequence_output[:, 1:]
+
+        linear_input = torch.cat([cls_token, patch_tokens.mean(dim=1)], dim=1)
+
+        logits = self.classifier(linear_input)
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    Dinov2 backbone, to be used with frameworks like DETR and MaskFormer.
+    """,
+    DINOV2_START_DOCSTRING,
+)
+class Dinov2Backbone(Dinov2PreTrainedModel, BackboneMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        super()._init_backbone(config)
+
+        self.num_features = [
+            config.hidden_size for _ in range(config.num_hidden_layers + 1)
+        ]
+        self.embeddings = Dinov2Embeddings(config)
+        self.encoder = Dinov2Encoder(config)
+
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> Dinov2PatchEmbeddings:
+        return self.embeddings.patch_embeddings
+
+    @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> BackboneOutput:
+        """
+        Returns:
+
+        Examples:
+
+        ```python
+        >>> from transformers import AutoImageProcessor, AutoBackbone
+        >>> import torch
+        >>> from PIL import Image
+        >>> import requests
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
+        >>> model = AutoBackbone.from_pretrained(
+        ...     "facebook/dinov2-base", out_features=["stage2", "stage5", "stage8", "stage11"]
+        ... )
+
+        >>> inputs = processor(image, return_tensors="pt")
+
+        >>> outputs = model(**inputs)
+        >>> feature_maps = outputs.feature_maps
+        >>> list(feature_maps[-1].shape)
+        [1, 768, 16, 16]
+        ```"""
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+
+        embedding_output = self.embeddings(pixel_values)
+
+        outputs = self.encoder(
+            embedding_output,
+            output_hidden_states=True,
+            output_attentions=output_attentions,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs.hidden_states if return_dict else outputs[1]
+
+        feature_maps = ()
+        for stage, hidden_state in zip(self.stage_names, hidden_states):
+            if stage in self.out_features:
+                if self.config.apply_layernorm:
+                    hidden_state = self.layernorm(hidden_state)
+                if self.config.reshape_hidden_states:
+                    batch_size, _, height, width = pixel_values.shape
+                    patch_size = self.config.patch_size
+                    hidden_state = hidden_state[:, 1:, :].reshape(
+                        batch_size, width // patch_size, height // patch_size, -1
+                    )
+                    hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()
+                feature_maps += (hidden_state,)
+
+        if not return_dict:
+            if output_hidden_states:
+                output = (feature_maps,) + outputs[1:]
+            else:
+                output = (feature_maps,) + outputs[2:]
+            return output
+
+        return BackboneOutput(
+            feature_maps=feature_maps,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            attentions=outputs.attentions if output_attentions else None,
+        )
+
+
+class CustomPatchEmbeddings(nn.Module):
+    """
+    This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+    `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
+    Transformer.
+    """
+
+    def __init__(
+        self, image_size: int, patch_size: int, num_channels: int, hidden_size: int
+    ):
+        super().__init__()
+
+        image_size = (
+            image_size
+            if isinstance(image_size, collections.abc.Iterable)
+            else (image_size, image_size)
+        )
+        patch_size = (
+            patch_size
+            if isinstance(patch_size, collections.abc.Iterable)
+            else (patch_size, patch_size)
+        )
+        num_patches = (image_size[1] // patch_size[1]) * (
+            image_size[0] // patch_size[0]
+        )
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.num_patches = num_patches
+
+        self.projection = nn.Conv2d(
+            num_channels, hidden_size, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        num_channels = pixel_values.shape[1]
+        if num_channels != self.num_channels:
+            raise ValueError(
+                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+                f" Expected {self.num_channels} but got {num_channels}."
+            )
+        embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
+        return embeddings
+
+
+class CustomEmbeddings(nn.Module):
+    """
+    Construct the CLS token, mask token, position and patch embeddings.
+    """
+
+    def __init__(
+        self, image_size: int, patch_size: int, num_channels: int, hidden_size: int
+    ) -> None:
+        super().__init__()
+
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.hidden_size = hidden_size
+
+        self.cls_token = nn.Parameter(torch.randn(1, 1, self.hidden_size))
+
+        self.patch_embeddings = CustomPatchEmbeddings(
+            image_size, patch_size, num_channels, hidden_size
+        )
+        num_patches = self.patch_embeddings.num_patches
+        self.position_embeddings = nn.Parameter(
+            torch.randn(1, num_patches + 1, self.hidden_size)
+        )
+
+    def interpolate_pos_encoding(
+        self, embeddings: torch.Tensor, height: int, width: int
+    ) -> torch.Tensor:
+        """
+        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
+        resolution images.
+
+        Source:
+        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
+        """
+
+        num_patches = embeddings.shape[1] - 1
+        num_positions = self.position_embeddings.shape[1] - 1
+        if num_patches == num_positions and height == width:
+            return self.position_embeddings
+        class_pos_embed = self.position_embeddings[:, 0]
+        patch_pos_embed = self.position_embeddings[:, 1:]
+        dim = embeddings.shape[-1]
+        height = height // self.patch_size
+        width = width // self.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        height, width = height + 0.1, width + 0.1
+        patch_pos_embed = patch_pos_embed.reshape(
+            1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim
+        )
+        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed,
+            scale_factor=(
+                height / math.sqrt(num_positions),
+                width / math.sqrt(num_positions),
+            ),
+            mode="bicubic",
+            align_corners=False,
+        )
+        if (
+            int(height) != patch_pos_embed.shape[-2]
+            or int(width) != patch_pos_embed.shape[-1]
+        ):
+            raise ValueError(
+                "Width or height does not match with the interpolated position embeddings"
+            )
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+    ) -> torch.Tensor:
+        batch_size, _, height, width = pixel_values.shape
+        patch_embeddings = self.patch_embeddings(pixel_values)
+        embeddings = patch_embeddings
+
+        # add the [CLS] token to the embedded patch tokens
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        embeddings = torch.cat((cls_tokens, embeddings), dim=1)
+
+        # add positional encoding to each token
+        embeddings = embeddings + self.interpolate_pos_encoding(
+            embeddings, height, width
+        )
+
+        return embeddings

spar3d/models/tokenizers/image.py

@@ -0,0 +1,99 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from jaxtyping import Float
+from torch import Tensor
+
+from spar3d.models.tokenizers.dinov2 import Dinov2Model
+from spar3d.models.transformers.attention import Modulation
+from spar3d.models.utils import BaseModule
+
+
+class DINOV2SingleImageTokenizer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        pretrained_model_name_or_path: str = "facebook/dinov2-large"
+        width: int = 512
+        height: int = 512
+        modulation_cond_dim: int = 768
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.model = Dinov2Model.from_pretrained(self.cfg.pretrained_model_name_or_path)
+
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+        self.model.eval()
+
+        self.model.set_gradient_checkpointing(False)
+
+        # add modulation
+        modulations = []
+        for layer in self.model.encoder.layer:
+            norm1_modulation = Modulation(
+                self.model.config.hidden_size,
+                self.cfg.modulation_cond_dim,
+                zero_init=True,
+                single_layer=True,
+            )
+            norm2_modulation = Modulation(
+                self.model.config.hidden_size,
+                self.cfg.modulation_cond_dim,
+                zero_init=True,
+                single_layer=True,
+            )
+            layer.register_ada_norm_modulation(norm1_modulation, norm2_modulation)
+            modulations += [norm1_modulation, norm2_modulation]
+        self.modulations = nn.ModuleList(modulations)
+
+        self.register_buffer(
+            "image_mean",
+            torch.as_tensor([0.485, 0.456, 0.406]).reshape(1, 1, 3, 1, 1),
+            persistent=False,
+        )
+        self.register_buffer(
+            "image_std",
+            torch.as_tensor([0.229, 0.224, 0.225]).reshape(1, 1, 3, 1, 1),
+            persistent=False,
+        )
+
+    def forward(
+        self,
+        images: Float[Tensor, "B *N C H W"],
+        modulation_cond: Optional[Float[Tensor, "B *N Cc"]],
+        **kwargs,
+    ) -> Float[Tensor, "B *N Ct Nt"]:
+        model = self.model
+
+        packed = False
+        if images.ndim == 4:
+            packed = True
+            images = images.unsqueeze(1)
+            if modulation_cond is not None:
+                assert modulation_cond.ndim == 2
+                modulation_cond = modulation_cond.unsqueeze(1)
+
+        batch_size, n_input_views = images.shape[:2]
+        images = (images - self.image_mean) / self.image_std
+        out = model(
+            rearrange(images, "B N C H W -> (B N) C H W"),
+            modulation_cond=rearrange(modulation_cond, "B N Cc -> (B N) Cc")
+            if modulation_cond is not None
+            else None,
+        )
+        local_features = out.last_hidden_state
+        local_features = local_features.permute(0, 2, 1)
+        local_features = rearrange(
+            local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size
+        )
+        if packed:
+            local_features = local_features.squeeze(1)
+
+        return local_features
+
+    def detokenize(self, *args, **kwargs):
+        raise NotImplementedError

spar3d/models/tokenizers/point.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from jaxtyping import Float
+from torch import Tensor
+
+from spar3d.models.transformers.transformer_1d import Transformer1D
+from spar3d.models.utils import BaseModule
+
+
+class TransformerPointTokenizer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        num_attention_heads: int = 16
+        attention_head_dim: int = 64
+        in_channels: Optional[int] = 6
+        out_channels: Optional[int] = 1024
+        num_layers: int = 16
+        norm_num_groups: int = 32
+        attention_bias: bool = False
+        activation_fn: str = "geglu"
+        norm_elementwise_affine: bool = True
+
+    cfg: Config
+
+    def configure(self) -> None:
+        transformer_cfg = dict(self.cfg.copy())
+        # remove the non-transformer configs
+        transformer_cfg["in_channels"] = (
+            self.cfg.num_attention_heads * self.cfg.attention_head_dim
+        )
+        self.model = Transformer1D(transformer_cfg)
+        self.linear_in = torch.nn.Linear(
+            self.cfg.in_channels, transformer_cfg["in_channels"]
+        )
+        self.linear_out = torch.nn.Linear(
+            transformer_cfg["in_channels"], self.cfg.out_channels
+        )
+
+    def forward(
+        self, points: Float[Tensor, "B N Ci"], **kwargs
+    ) -> Float[Tensor, "B N Cp"]:
+        assert points.ndim == 3
+        inputs = self.linear_in(points).permute(0, 2, 1)  # B N Ci -> B Ci N
+        out = self.model(inputs).permute(0, 2, 1)  # B Ci N -> B N Ci
+        out = self.linear_out(out)  # B N Ci -> B N Co
+        return out
+
+    def detokenize(self, *args, **kwargs):
+        raise NotImplementedError

spar3d/models/tokenizers/triplane.py

@@ -0,0 +1,49 @@
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from jaxtyping import Float
+from torch import Tensor
+
+from spar3d.models.utils import BaseModule
+
+
+class TriplaneLearnablePositionalEmbedding(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        plane_size: int = 96
+        num_channels: int = 1024
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.embeddings = nn.Parameter(
+            torch.randn(
+                (3, self.cfg.num_channels, self.cfg.plane_size, self.cfg.plane_size),
+                dtype=torch.float32,
+            )
+            * 1
+            / math.sqrt(self.cfg.num_channels)
+        )
+
+    def forward(self, batch_size: int) -> Float[Tensor, "B Ct Nt"]:
+        return rearrange(
+            repeat(self.embeddings, "Np Ct Hp Wp -> B Np Ct Hp Wp", B=batch_size),
+            "B Np Ct Hp Wp -> B Ct (Np Hp Wp)",
+        )
+
+    def detokenize(
+        self, tokens: Float[Tensor, "B Ct Nt"]
+    ) -> Float[Tensor, "B 3 Ct Hp Wp"]:
+        batch_size, Ct, Nt = tokens.shape
+        assert Nt == self.cfg.plane_size**2 * 3
+        assert Ct == self.cfg.num_channels
+        return rearrange(
+            tokens,
+            "B Ct (Np Hp Wp) -> B Np Ct Hp Wp",
+            Np=3,
+            Hp=self.cfg.plane_size,
+            Wp=self.cfg.plane_size,
+        )