chores: rebase commits

threestudio/models/__init__.py

@@ -0,0 +1,9 @@
+from . import (
+    background,
+    exporters,
+    geometry,
+    guidance,
+    materials,
+    prompt_processors,
+    renderers,
+)

threestudio/models/background/__init__.py

@@ -0,0 +1,6 @@
+from . import (
+    base,
+    neural_environment_map_background,
+    solid_color_background,
+    textured_background,
+)

threestudio/models/background/base.py

@@ -0,0 +1,24 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.utils.base import BaseModule
+from threestudio.utils.typing import *
+
+
+class BaseBackground(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        pass
+
+    cfg: Config
+
+    def configure(self):
+        pass
+
+    def forward(self, dirs: Float[Tensor, "B H W 3"]) -> Float[Tensor, "B H W Nc"]:
+        raise NotImplementedError

threestudio/models/background/neural_environment_map_background.py

@@ -0,0 +1,71 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.ops import get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("neural-environment-map-background")
+class NeuralEnvironmentMapBackground(BaseBackground):
+    @dataclass
+    class Config(BaseBackground.Config):
+        n_output_dims: int = 3
+        color_activation: str = "sigmoid"
+        dir_encoding_config: dict = field(
+            default_factory=lambda: {"otype": "SphericalHarmonics", "degree": 3}
+        )
+        mlp_network_config: dict = field(
+            default_factory=lambda: {
+                "otype": "VanillaMLP",
+                "activation": "ReLU",
+                "n_neurons": 16,
+                "n_hidden_layers": 2,
+            }
+        )
+        random_aug: bool = False
+        random_aug_prob: float = 0.5
+        eval_color: Optional[Tuple[float, float, float]] = None
+
+        # multi-view diffusion
+        share_aug_bg: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.encoding = get_encoding(3, self.cfg.dir_encoding_config)
+        self.network = get_mlp(
+            self.encoding.n_output_dims,
+            self.cfg.n_output_dims,
+            self.cfg.mlp_network_config,
+        )
+
+    def forward(self, dirs: Float[Tensor, "B H W 3"]) -> Float[Tensor, "B H W Nc"]:
+        if not self.training and self.cfg.eval_color is not None:
+            return torch.ones(*dirs.shape[:-1], self.cfg.n_output_dims).to(
+                dirs
+            ) * torch.as_tensor(self.cfg.eval_color).to(dirs)
+        # viewdirs must be normalized before passing to this function
+        dirs = (dirs + 1.0) / 2.0  # (-1, 1) => (0, 1)
+        dirs_embd = self.encoding(dirs.view(-1, 3))
+        color = self.network(dirs_embd).view(*dirs.shape[:-1], self.cfg.n_output_dims)
+        color = get_activation(self.cfg.color_activation)(color)
+        if (
+            self.training
+            and self.cfg.random_aug
+            and random.random() < self.cfg.random_aug_prob
+        ):
+            # use random background color with probability random_aug_prob
+            n_color = 1 if self.cfg.share_aug_bg else dirs.shape[0]
+            color = color * 0 + (  # prevent checking for unused parameters in DDP
+                torch.rand(n_color, 1, 1, self.cfg.n_output_dims)
+                .to(dirs)
+                .expand(*dirs.shape[:-1], -1)
+            )
+        return color

threestudio/models/background/solid_color_background.py

@@ -0,0 +1,51 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.utils.typing import *
+
+
+@threestudio.register("solid-color-background")
+class SolidColorBackground(BaseBackground):
+    @dataclass
+    class Config(BaseBackground.Config):
+        n_output_dims: int = 3
+        color: Tuple = (1.0, 1.0, 1.0)
+        learned: bool = False
+        random_aug: bool = False
+        random_aug_prob: float = 0.5
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.env_color: Float[Tensor, "Nc"]
+        if self.cfg.learned:
+            self.env_color = nn.Parameter(
+                torch.as_tensor(self.cfg.color, dtype=torch.float32)
+            )
+        else:
+            self.register_buffer(
+                "env_color", torch.as_tensor(self.cfg.color, dtype=torch.float32)
+            )
+
+    def forward(self, dirs: Float[Tensor, "B H W 3"]) -> Float[Tensor, "B H W Nc"]:
+        color = torch.ones(*dirs.shape[:-1], self.cfg.n_output_dims).to(
+            dirs
+        ) * self.env_color.to(dirs)
+        if (
+            self.training
+            and self.cfg.random_aug
+            and random.random() < self.cfg.random_aug_prob
+        ):
+            # use random background color with probability random_aug_prob
+            color = color * 0 + (  # prevent checking for unused parameters in DDP
+                torch.rand(dirs.shape[0], 1, 1, self.cfg.n_output_dims)
+                .to(dirs)
+                .expand(*dirs.shape[:-1], -1)
+            )
+        return color

threestudio/models/background/textured_background.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.utils.ops import get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("textured-background")
+class TexturedBackground(BaseBackground):
+    @dataclass
+    class Config(BaseBackground.Config):
+        n_output_dims: int = 3
+        height: int = 64
+        width: int = 64
+        color_activation: str = "sigmoid"
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.texture = nn.Parameter(
+            torch.randn((1, self.cfg.n_output_dims, self.cfg.height, self.cfg.width))
+        )
+
+    def spherical_xyz_to_uv(self, dirs: Float[Tensor, "*B 3"]) -> Float[Tensor, "*B 2"]:
+        x, y, z = dirs[..., 0], dirs[..., 1], dirs[..., 2]
+        xy = (x**2 + y**2) ** 0.5
+        u = torch.atan2(xy, z) / torch.pi
+        v = torch.atan2(y, x) / (torch.pi * 2) + 0.5
+        uv = torch.stack([u, v], -1)
+        return uv
+
+    def forward(self, dirs: Float[Tensor, "*B 3"]) -> Float[Tensor, "*B Nc"]:
+        dirs_shape = dirs.shape[:-1]
+        uv = self.spherical_xyz_to_uv(dirs.reshape(-1, dirs.shape[-1]))
+        uv = 2 * uv - 1  # rescale to [-1, 1] for grid_sample
+        uv = uv.reshape(1, -1, 1, 2)
+        color = (
+            F.grid_sample(
+                self.texture,
+                uv,
+                mode="bilinear",
+                padding_mode="reflection",
+                align_corners=False,
+            )
+            .reshape(self.cfg.n_output_dims, -1)
+            .T.reshape(*dirs_shape, self.cfg.n_output_dims)
+        )
+        color = get_activation(self.cfg.color_activation)(color)
+        return color

threestudio/models/estimators.py

@@ -0,0 +1,118 @@
+from typing import Callable, List, Optional, Tuple
+
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
+
+import torch
+from nerfacc.data_specs import RayIntervals
+from nerfacc.estimators.base import AbstractEstimator
+from nerfacc.pdf import importance_sampling, searchsorted
+from nerfacc.volrend import render_transmittance_from_density
+from torch import Tensor
+
+
+class ImportanceEstimator(AbstractEstimator):
+    def __init__(
+        self,
+    ) -> None:
+        super().__init__()
+
+    @torch.no_grad()
+    def sampling(
+        self,
+        prop_sigma_fns: List[Callable],
+        prop_samples: List[int],
+        num_samples: int,
+        # rendering options
+        n_rays: int,
+        near_plane: float,
+        far_plane: float,
+        sampling_type: Literal["uniform", "lindisp"] = "uniform",
+        # training options
+        stratified: bool = False,
+        requires_grad: bool = False,
+    ) -> Tuple[Tensor, Tensor]:
+        """Sampling with CDFs from proposal networks.
+
+        Args:
+            prop_sigma_fns: Proposal network evaluate functions. It should be a list
+                of functions that take in samples {t_starts (n_rays, n_samples),
+                t_ends (n_rays, n_samples)} and returns the post-activation densities
+                (n_rays, n_samples).
+            prop_samples: Number of samples to draw from each proposal network. Should
+                be the same length as `prop_sigma_fns`.
+            num_samples: Number of samples to draw in the end.
+            n_rays: Number of rays.
+            near_plane: Near plane.
+            far_plane: Far plane.
+            sampling_type: Sampling type. Either "uniform" or "lindisp". Default to
+                "lindisp".
+            stratified: Whether to use stratified sampling. Default to `False`.
+
+        Returns:
+            A tuple of {Tensor, Tensor}:
+
+            - **t_starts**: The starts of the samples. Shape (n_rays, num_samples).
+            - **t_ends**: The ends of the samples. Shape (n_rays, num_samples).
+
+        """
+        assert len(prop_sigma_fns) == len(prop_samples), (
+            "The number of proposal networks and the number of samples "
+            "should be the same."
+        )
+        cdfs = torch.cat(
+            [
+                torch.zeros((n_rays, 1), device=self.device),
+                torch.ones((n_rays, 1), device=self.device),
+            ],
+            dim=-1,
+        )
+        intervals = RayIntervals(vals=cdfs)
+
+        for level_fn, level_samples in zip(prop_sigma_fns, prop_samples):
+            intervals, _ = importance_sampling(
+                intervals, cdfs, level_samples, stratified
+            )
+            t_vals = _transform_stot(
+                sampling_type, intervals.vals, near_plane, far_plane
+            )
+            t_starts = t_vals[..., :-1]
+            t_ends = t_vals[..., 1:]
+
+            with torch.set_grad_enabled(requires_grad):
+                sigmas = level_fn(t_starts, t_ends)
+                assert sigmas.shape == t_starts.shape
+                trans, _ = render_transmittance_from_density(t_starts, t_ends, sigmas)
+                cdfs = 1.0 - torch.cat([trans, torch.zeros_like(trans[:, :1])], dim=-1)
+
+        intervals, _ = importance_sampling(intervals, cdfs, num_samples, stratified)
+        t_vals_fine = _transform_stot(
+            sampling_type, intervals.vals, near_plane, far_plane
+        )
+
+        t_vals = torch.cat([t_vals, t_vals_fine], dim=-1)
+        t_vals, _ = torch.sort(t_vals, dim=-1)
+
+        t_starts_ = t_vals[..., :-1]
+        t_ends_ = t_vals[..., 1:]
+
+        return t_starts_, t_ends_
+
+
+def _transform_stot(
+    transform_type: Literal["uniform", "lindisp"],
+    s_vals: torch.Tensor,
+    t_min: torch.Tensor,
+    t_max: torch.Tensor,
+) -> torch.Tensor:
+    if transform_type == "uniform":
+        _contract_fn, _icontract_fn = lambda x: x, lambda x: x
+    elif transform_type == "lindisp":
+        _contract_fn, _icontract_fn = lambda x: 1 / x, lambda x: 1 / x
+    else:
+        raise ValueError(f"Unknown transform_type: {transform_type}")
+    s_min, s_max = _contract_fn(t_min), _contract_fn(t_max)
+    icontract_fn = lambda s: _icontract_fn(s * s_max + (1 - s) * s_min)
+    return icontract_fn(s_vals)

threestudio/models/exporters/__init__.py

@@ -0,0 +1 @@
+from . import base, mesh_exporter

threestudio/models/exporters/base.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.geometry.base import BaseImplicitGeometry
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.utils.base import BaseObject
+from threestudio.utils.typing import *
+
+
+@dataclass
+class ExporterOutput:
+    save_name: str
+    save_type: str
+    params: Dict[str, Any]
+
+
+class Exporter(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        save_video: bool = False
+
+    cfg: Config
+
+    def configure(
+        self,
+        geometry: BaseImplicitGeometry,
+        material: BaseMaterial,
+        background: BaseBackground,
+    ) -> None:
+        @dataclass
+        class SubModules:
+            geometry: BaseImplicitGeometry
+            material: BaseMaterial
+            background: BaseBackground
+
+        self.sub_modules = SubModules(geometry, material, background)
+
+    @property
+    def geometry(self) -> BaseImplicitGeometry:
+        return self.sub_modules.geometry
+
+    @property
+    def material(self) -> BaseMaterial:
+        return self.sub_modules.material
+
+    @property
+    def background(self) -> BaseBackground:
+        return self.sub_modules.background
+
+    def __call__(self, *args, **kwargs) -> List[ExporterOutput]:
+        raise NotImplementedError
+
+
+@threestudio.register("dummy-exporter")
+class DummyExporter(Exporter):
+    def __call__(self, *args, **kwargs) -> List[ExporterOutput]:
+        # DummyExporter does not export anything
+        return []

threestudio/models/exporters/mesh_exporter.py

@@ -0,0 +1,175 @@
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+import torch
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.exporters.base import Exporter, ExporterOutput
+from threestudio.models.geometry.base import BaseImplicitGeometry
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.mesh import Mesh
+from threestudio.utils.rasterize import NVDiffRasterizerContext
+from threestudio.utils.typing import *
+
+
+@threestudio.register("mesh-exporter")
+class MeshExporter(Exporter):
+    @dataclass
+    class Config(Exporter.Config):
+        fmt: str = "obj-mtl"  # in ['obj-mtl', 'obj'], TODO: fbx
+        save_name: str = "model"
+        save_normal: bool = False
+        save_uv: bool = True
+        save_texture: bool = True
+        texture_size: int = 1024
+        texture_format: str = "jpg"
+        xatlas_chart_options: dict = field(default_factory=dict)
+        xatlas_pack_options: dict = field(default_factory=dict)
+        context_type: str = "gl"
+
+    cfg: Config
+
+    def configure(
+        self,
+        geometry: BaseImplicitGeometry,
+        material: BaseMaterial,
+        background: BaseBackground,
+    ) -> None:
+        super().configure(geometry, material, background)
+        self.ctx = NVDiffRasterizerContext(self.cfg.context_type, self.device)
+
+    def __call__(self) -> List[ExporterOutput]:
+        mesh: Mesh = self.geometry.isosurface()
+
+        if self.cfg.fmt == "obj-mtl":
+            return self.export_obj_with_mtl(mesh)
+        elif self.cfg.fmt == "obj":
+            return self.export_obj(mesh)
+        else:
+            raise ValueError(f"Unsupported mesh export format: {self.cfg.fmt}")
+
+    def export_obj_with_mtl(self, mesh: Mesh) -> List[ExporterOutput]:
+        params = {
+            "mesh": mesh,
+            "save_mat": True,
+            "save_normal": self.cfg.save_normal,
+            "save_uv": self.cfg.save_uv,
+            "save_vertex_color": False,
+            "map_Kd": None,  # Base Color
+            "map_Ks": None,  # Specular
+            "map_Bump": None,  # Normal
+            # ref: https://en.wikipedia.org/wiki/Wavefront_.obj_file#Physically-based_Rendering
+            "map_Pm": None,  # Metallic
+            "map_Pr": None,  # Roughness
+            "map_format": self.cfg.texture_format,
+        }
+
+        if self.cfg.save_uv:
+            mesh.unwrap_uv(self.cfg.xatlas_chart_options, self.cfg.xatlas_pack_options)
+
+        if self.cfg.save_texture:
+            threestudio.info("Exporting textures ...")
+            assert self.cfg.save_uv, "save_uv must be True when save_texture is True"
+            # clip space transform
+            uv_clip = mesh.v_tex * 2.0 - 1.0
+            # pad to four component coordinate
+            uv_clip4 = torch.cat(
+                (
+                    uv_clip,
+                    torch.zeros_like(uv_clip[..., 0:1]),
+                    torch.ones_like(uv_clip[..., 0:1]),
+                ),
+                dim=-1,
+            )
+            # rasterize
+            rast, _ = self.ctx.rasterize_one(
+                uv_clip4, mesh.t_tex_idx, (self.cfg.texture_size, self.cfg.texture_size)
+            )
+
+            hole_mask = ~(rast[:, :, 3] > 0)
+
+            def uv_padding(image):
+                uv_padding_size = self.cfg.xatlas_pack_options.get("padding", 2)
+                inpaint_image = (
+                    cv2.inpaint(
+                        (image.detach().cpu().numpy() * 255).astype(np.uint8),
+                        (hole_mask.detach().cpu().numpy() * 255).astype(np.uint8),
+                        uv_padding_size,
+                        cv2.INPAINT_TELEA,
+                    )
+                    / 255.0
+                )
+                return torch.from_numpy(inpaint_image).to(image)
+
+            # Interpolate world space position
+            gb_pos, _ = self.ctx.interpolate_one(
+                mesh.v_pos, rast[None, ...], mesh.t_pos_idx
+            )
+            gb_pos = gb_pos[0]
+
+            # Sample out textures from MLP
+            geo_out = self.geometry.export(points=gb_pos)
+            mat_out = self.material.export(points=gb_pos, **geo_out)
+
+            threestudio.info(
+                "Perform UV padding on texture maps to avoid seams, may take a while ..."
+            )
+
+            if "albedo" in mat_out:
+                params["map_Kd"] = uv_padding(mat_out["albedo"])
+            else:
+                threestudio.warn(
+                    "save_texture is True but no albedo texture found, using default white texture"
+                )
+            if "metallic" in mat_out:
+                params["map_Pm"] = uv_padding(mat_out["metallic"])
+            if "roughness" in mat_out:
+                params["map_Pr"] = uv_padding(mat_out["roughness"])
+            if "bump" in mat_out:
+                params["map_Bump"] = uv_padding(mat_out["bump"])
+            # TODO: map_Ks
+        return [
+            ExporterOutput(
+                save_name=f"{self.cfg.save_name}.obj", save_type="obj", params=params
+            )
+        ]
+
+    def export_obj(self, mesh: Mesh) -> List[ExporterOutput]:
+        params = {
+            "mesh": mesh,
+            "save_mat": False,
+            "save_normal": self.cfg.save_normal,
+            "save_uv": self.cfg.save_uv,
+            "save_vertex_color": False,
+            "map_Kd": None,  # Base Color
+            "map_Ks": None,  # Specular
+            "map_Bump": None,  # Normal
+            # ref: https://en.wikipedia.org/wiki/Wavefront_.obj_file#Physically-based_Rendering
+            "map_Pm": None,  # Metallic
+            "map_Pr": None,  # Roughness
+            "map_format": self.cfg.texture_format,
+        }
+
+        if self.cfg.save_uv:
+            mesh.unwrap_uv(self.cfg.xatlas_chart_options, self.cfg.xatlas_pack_options)
+
+        if self.cfg.save_texture:
+            threestudio.info("Exporting textures ...")
+            geo_out = self.geometry.export(points=mesh.v_pos)
+            mat_out = self.material.export(points=mesh.v_pos, **geo_out)
+
+            if "albedo" in mat_out:
+                mesh.set_vertex_color(mat_out["albedo"])
+                params["save_vertex_color"] = True
+            else:
+                threestudio.warn(
+                    "save_texture is True but no albedo texture found, not saving vertex color"
+                )
+
+        return [
+            ExporterOutput(
+                save_name=f"{self.cfg.save_name}.obj", save_type="obj", params=params
+            )
+        ]

threestudio/models/geometry/__init__.py

@@ -0,0 +1,8 @@
+from . import (
+    base,
+    custom_mesh,
+    implicit_sdf,
+    implicit_volume,
+    tetrahedra_sdf_grid,
+    volume_grid,
+)

threestudio/models/geometry/base.py

@@ -0,0 +1,209 @@
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.isosurface import (
+    IsosurfaceHelper,
+    MarchingCubeCPUHelper,
+    MarchingTetrahedraHelper,
+)
+from threestudio.models.mesh import Mesh
+from threestudio.utils.base import BaseModule
+from threestudio.utils.ops import chunk_batch, scale_tensor
+from threestudio.utils.typing import *
+
+
+def contract_to_unisphere(
+    x: Float[Tensor, "... 3"], bbox: Float[Tensor, "2 3"], unbounded: bool = False
+) -> Float[Tensor, "... 3"]:
+    if unbounded:
+        x = scale_tensor(x, bbox, (0, 1))
+        x = x * 2 - 1  # aabb is at [-1, 1]
+        mag = x.norm(dim=-1, keepdim=True)
+        mask = mag.squeeze(-1) > 1
+        x[mask] = (2 - 1 / mag[mask]) * (x[mask] / mag[mask])
+        x = x / 4 + 0.5  # [-inf, inf] is at [0, 1]
+    else:
+        x = scale_tensor(x, bbox, (0, 1))
+    return x
+
+
+class BaseGeometry(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        pass
+
+    cfg: Config
+
+    @staticmethod
+    def create_from(
+        other: "BaseGeometry", cfg: Optional[Union[dict, DictConfig]] = None, **kwargs
+    ) -> "BaseGeometry":
+        raise TypeError(
+            f"Cannot create {BaseGeometry.__name__} from {other.__class__.__name__}"
+        )
+
+    def export(self, *args, **kwargs) -> Dict[str, Any]:
+        return {}
+
+
+class BaseImplicitGeometry(BaseGeometry):
+    @dataclass
+    class Config(BaseGeometry.Config):
+        radius: float = 1.0
+        isosurface: bool = True
+        isosurface_method: str = "mt"
+        isosurface_resolution: int = 128
+        isosurface_threshold: Union[float, str] = 0.0
+        isosurface_chunk: int = 0
+        isosurface_coarse_to_fine: bool = True
+        isosurface_deformable_grid: bool = False
+        isosurface_remove_outliers: bool = True
+        isosurface_outlier_n_faces_threshold: Union[int, float] = 0.01
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.bbox: Float[Tensor, "2 3"]
+        self.register_buffer(
+            "bbox",
+            torch.as_tensor(
+                [
+                    [-self.cfg.radius, -self.cfg.radius, -self.cfg.radius],
+                    [self.cfg.radius, self.cfg.radius, self.cfg.radius],
+                ],
+                dtype=torch.float32,
+            ),
+        )
+        self.isosurface_helper: Optional[IsosurfaceHelper] = None
+        self.unbounded: bool = False
+
+    def _initilize_isosurface_helper(self):
+        if self.cfg.isosurface and self.isosurface_helper is None:
+            if self.cfg.isosurface_method == "mc-cpu":
+                self.isosurface_helper = MarchingCubeCPUHelper(
+                    self.cfg.isosurface_resolution
+                ).to(self.device)
+            elif self.cfg.isosurface_method == "mt":
+                self.isosurface_helper = MarchingTetrahedraHelper(
+                    self.cfg.isosurface_resolution,
+                    f"load/tets/{self.cfg.isosurface_resolution}_tets.npz",
+                ).to(self.device)
+            else:
+                raise AttributeError(
+                    "Unknown isosurface method {self.cfg.isosurface_method}"
+                )
+
+    def forward(
+        self, points: Float[Tensor, "*N Di"], output_normal: bool = False
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        raise NotImplementedError
+
+    def forward_field(
+        self, points: Float[Tensor, "*N Di"]
+    ) -> Tuple[Float[Tensor, "*N 1"], Optional[Float[Tensor, "*N 3"]]]:
+        # return the value of the implicit field, could be density / signed distance
+        # also return a deformation field if the grid vertices can be optimized
+        raise NotImplementedError
+
+    def forward_level(
+        self, field: Float[Tensor, "*N 1"], threshold: float
+    ) -> Float[Tensor, "*N 1"]:
+        # return the value of the implicit field, where the zero level set represents the surface
+        raise NotImplementedError
+
+    def _isosurface(self, bbox: Float[Tensor, "2 3"], fine_stage: bool = False) -> Mesh:
+        def batch_func(x):
+            # scale to bbox as the input vertices are in [0, 1]
+            field, deformation = self.forward_field(
+                scale_tensor(
+                    x.to(bbox.device), self.isosurface_helper.points_range, bbox
+                ),
+            )
+            field = field.to(
+                x.device
+            )  # move to the same device as the input (could be CPU)
+            if deformation is not None:
+                deformation = deformation.to(x.device)
+            return field, deformation
+
+        assert self.isosurface_helper is not None
+
+        field, deformation = chunk_batch(
+            batch_func,
+            self.cfg.isosurface_chunk,
+            self.isosurface_helper.grid_vertices,
+        )
+
+        threshold: float
+
+        if isinstance(self.cfg.isosurface_threshold, float):
+            threshold = self.cfg.isosurface_threshold
+        elif self.cfg.isosurface_threshold == "auto":
+            eps = 1.0e-5
+            threshold = field[field > eps].mean().item()
+            threestudio.info(
+                f"Automatically determined isosurface threshold: {threshold}"
+            )
+        else:
+            raise TypeError(
+                f"Unknown isosurface_threshold {self.cfg.isosurface_threshold}"
+            )
+
+        level = self.forward_level(field, threshold)
+        mesh: Mesh = self.isosurface_helper(level, deformation=deformation)
+        mesh.v_pos = scale_tensor(
+            mesh.v_pos, self.isosurface_helper.points_range, bbox
+        )  # scale to bbox as the grid vertices are in [0, 1]
+        mesh.add_extra("bbox", bbox)
+
+        if self.cfg.isosurface_remove_outliers:
+            # remove outliers components with small number of faces
+            # only enabled when the mesh is not differentiable
+            mesh = mesh.remove_outlier(self.cfg.isosurface_outlier_n_faces_threshold)
+
+        return mesh
+
+    def isosurface(self) -> Mesh:
+        if not self.cfg.isosurface:
+            raise NotImplementedError(
+                "Isosurface is not enabled in the current configuration"
+            )
+        self._initilize_isosurface_helper()
+        if self.cfg.isosurface_coarse_to_fine:
+            threestudio.debug("First run isosurface to get a tight bounding box ...")
+            with torch.no_grad():
+                mesh_coarse = self._isosurface(self.bbox)
+            vmin, vmax = mesh_coarse.v_pos.amin(dim=0), mesh_coarse.v_pos.amax(dim=0)
+            vmin_ = (vmin - (vmax - vmin) * 0.1).max(self.bbox[0])
+            vmax_ = (vmax + (vmax - vmin) * 0.1).min(self.bbox[1])
+            threestudio.debug("Run isosurface again with the tight bounding box ...")
+            mesh = self._isosurface(torch.stack([vmin_, vmax_], dim=0), fine_stage=True)
+        else:
+            mesh = self._isosurface(self.bbox)
+        return mesh
+
+
+class BaseExplicitGeometry(BaseGeometry):
+    @dataclass
+    class Config(BaseGeometry.Config):
+        radius: float = 1.0
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.bbox: Float[Tensor, "2 3"]
+        self.register_buffer(
+            "bbox",
+            torch.as_tensor(
+                [
+                    [-self.cfg.radius, -self.cfg.radius, -self.cfg.radius],
+                    [self.cfg.radius, self.cfg.radius, self.cfg.radius],
+                ],
+                dtype=torch.float32,
+            ),
+        )

threestudio/models/geometry/custom_mesh.py

@@ -0,0 +1,178 @@
+import os
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.geometry.base import (
+    BaseExplicitGeometry,
+    BaseGeometry,
+    contract_to_unisphere,
+)
+from threestudio.models.mesh import Mesh
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.ops import scale_tensor
+from threestudio.utils.typing import *
+
+
+@threestudio.register("custom-mesh")
+class CustomMesh(BaseExplicitGeometry):
+    @dataclass
+    class Config(BaseExplicitGeometry.Config):
+        n_input_dims: int = 3
+        n_feature_dims: int = 3
+        pos_encoding_config: dict = field(
+            default_factory=lambda: {
+                "otype": "HashGrid",
+                "n_levels": 16,
+                "n_features_per_level": 2,
+                "log2_hashmap_size": 19,
+                "base_resolution": 16,
+                "per_level_scale": 1.447269237440378,
+            }
+        )
+        mlp_network_config: dict = field(
+            default_factory=lambda: {
+                "otype": "VanillaMLP",
+                "activation": "ReLU",
+                "output_activation": "none",
+                "n_neurons": 64,
+                "n_hidden_layers": 1,
+            }
+        )
+        shape_init: str = ""
+        shape_init_params: Optional[Any] = None
+        shape_init_mesh_up: str = "+z"
+        shape_init_mesh_front: str = "+x"
+
+    cfg: Config
+
+    def configure(self) -> None:
+        super().configure()
+
+        self.encoding = get_encoding(
+            self.cfg.n_input_dims, self.cfg.pos_encoding_config
+        )
+        self.feature_network = get_mlp(
+            self.encoding.n_output_dims,
+            self.cfg.n_feature_dims,
+            self.cfg.mlp_network_config,
+        )
+
+        # Initialize custom mesh
+        if self.cfg.shape_init.startswith("mesh:"):
+            assert isinstance(self.cfg.shape_init_params, float)
+            mesh_path = self.cfg.shape_init[5:]
+            if not os.path.exists(mesh_path):
+                raise ValueError(f"Mesh file {mesh_path} does not exist.")
+
+            import trimesh
+
+            scene = trimesh.load(mesh_path)
+            if isinstance(scene, trimesh.Trimesh):
+                mesh = scene
+            elif isinstance(scene, trimesh.scene.Scene):
+                mesh = trimesh.Trimesh()
+                for obj in scene.geometry.values():
+                    mesh = trimesh.util.concatenate([mesh, obj])
+            else:
+                raise ValueError(f"Unknown mesh type at {mesh_path}.")
+
+            # move to center
+            centroid = mesh.vertices.mean(0)
+            mesh.vertices = mesh.vertices - centroid
+
+            # align to up-z and front-x
+            dirs = ["+x", "+y", "+z", "-x", "-y", "-z"]
+            dir2vec = {
+                "+x": np.array([1, 0, 0]),
+                "+y": np.array([0, 1, 0]),
+                "+z": np.array([0, 0, 1]),
+                "-x": np.array([-1, 0, 0]),
+                "-y": np.array([0, -1, 0]),
+                "-z": np.array([0, 0, -1]),
+            }
+            if (
+                self.cfg.shape_init_mesh_up not in dirs
+                or self.cfg.shape_init_mesh_front not in dirs
+            ):
+                raise ValueError(
+                    f"shape_init_mesh_up and shape_init_mesh_front must be one of {dirs}."
+                )
+            if self.cfg.shape_init_mesh_up[1] == self.cfg.shape_init_mesh_front[1]:
+                raise ValueError(
+                    "shape_init_mesh_up and shape_init_mesh_front must be orthogonal."
+                )
+            z_, x_ = (
+                dir2vec[self.cfg.shape_init_mesh_up],
+                dir2vec[self.cfg.shape_init_mesh_front],
+            )
+            y_ = np.cross(z_, x_)
+            std2mesh = np.stack([x_, y_, z_], axis=0).T
+            mesh2std = np.linalg.inv(std2mesh)
+
+            # scaling
+            scale = np.abs(mesh.vertices).max()
+            mesh.vertices = mesh.vertices / scale * self.cfg.shape_init_params
+            mesh.vertices = np.dot(mesh2std, mesh.vertices.T).T
+
+            v_pos = torch.tensor(mesh.vertices, dtype=torch.float32).to(self.device)
+            t_pos_idx = torch.tensor(mesh.faces, dtype=torch.int64).to(self.device)
+            self.mesh = Mesh(v_pos=v_pos, t_pos_idx=t_pos_idx)
+            self.register_buffer(
+                "v_buffer",
+                v_pos,
+            )
+            self.register_buffer(
+                "t_buffer",
+                t_pos_idx,
+            )
+
+        else:
+            raise ValueError(
+                f"Unknown shape initialization type: {self.cfg.shape_init}"
+            )
+        print(self.mesh.v_pos.device)
+
+    def isosurface(self) -> Mesh:
+        if hasattr(self, "mesh"):
+            return self.mesh
+        elif hasattr(self, "v_buffer"):
+            self.mesh = Mesh(v_pos=self.v_buffer, t_pos_idx=self.t_buffer)
+            return self.mesh
+        else:
+            raise ValueError(f"custom mesh is not initialized")
+
+    def forward(
+        self, points: Float[Tensor, "*N Di"], output_normal: bool = False
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        assert (
+            output_normal == False
+        ), f"Normal output is not supported for {self.__class__.__name__}"
+        points_unscaled = points  # points in the original scale
+        points = contract_to_unisphere(points, self.bbox)  # points normalized to (0, 1)
+        enc = self.encoding(points.view(-1, self.cfg.n_input_dims))
+        features = self.feature_network(enc).view(
+            *points.shape[:-1], self.cfg.n_feature_dims
+        )
+        return {"features": features}
+
+    def export(self, points: Float[Tensor, "*N Di"], **kwargs) -> Dict[str, Any]:
+        out: Dict[str, Any] = {}
+        if self.cfg.n_feature_dims == 0:
+            return out
+        points_unscaled = points
+        points = contract_to_unisphere(points_unscaled, self.bbox)
+        enc = self.encoding(points.reshape(-1, self.cfg.n_input_dims))
+        features = self.feature_network(enc).view(
+            *points.shape[:-1], self.cfg.n_feature_dims
+        )
+        out.update(
+            {
+                "features": features,
+            }
+        )
+        return out

threestudio/models/geometry/implicit_sdf.py

@@ -0,0 +1,413 @@
+import os
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.geometry.base import BaseImplicitGeometry, contract_to_unisphere
+from threestudio.models.mesh import Mesh
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.misc import broadcast, get_rank
+from threestudio.utils.typing import *
+
+
+@threestudio.register("implicit-sdf")
+class ImplicitSDF(BaseImplicitGeometry):
+    @dataclass
+    class Config(BaseImplicitGeometry.Config):
+        n_input_dims: int = 3
+        n_feature_dims: int = 3
+        pos_encoding_config: dict = field(
+            default_factory=lambda: {
+                "otype": "HashGrid",
+                "n_levels": 16,
+                "n_features_per_level": 2,
+                "log2_hashmap_size": 19,
+                "base_resolution": 16,
+                "per_level_scale": 1.447269237440378,
+            }
+        )
+        mlp_network_config: dict = field(
+            default_factory=lambda: {
+                "otype": "VanillaMLP",
+                "activation": "ReLU",
+                "output_activation": "none",
+                "n_neurons": 64,
+                "n_hidden_layers": 1,
+            }
+        )
+        normal_type: Optional[
+            str
+        ] = "finite_difference"  # in ['pred', 'finite_difference', 'finite_difference_laplacian']
+        finite_difference_normal_eps: Union[
+            float, str
+        ] = 0.01  # in [float, "progressive"]
+        shape_init: Optional[str] = None
+        shape_init_params: Optional[Any] = None
+        shape_init_mesh_up: str = "+z"
+        shape_init_mesh_front: str = "+x"
+        force_shape_init: bool = False
+        sdf_bias: Union[float, str] = 0.0
+        sdf_bias_params: Optional[Any] = None
+
+        # no need to removal outlier for SDF
+        isosurface_remove_outliers: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        super().configure()
+        self.encoding = get_encoding(
+            self.cfg.n_input_dims, self.cfg.pos_encoding_config
+        )
+        self.sdf_network = get_mlp(
+            self.encoding.n_output_dims, 1, self.cfg.mlp_network_config
+        )
+
+        if self.cfg.n_feature_dims > 0:
+            self.feature_network = get_mlp(
+                self.encoding.n_output_dims,
+                self.cfg.n_feature_dims,
+                self.cfg.mlp_network_config,
+            )
+
+        if self.cfg.normal_type == "pred":
+            self.normal_network = get_mlp(
+                self.encoding.n_output_dims, 3, self.cfg.mlp_network_config
+            )
+        if self.cfg.isosurface_deformable_grid:
+            assert (
+                self.cfg.isosurface_method == "mt"
+            ), "isosurface_deformable_grid only works with mt"
+            self.deformation_network = get_mlp(
+                self.encoding.n_output_dims, 3, self.cfg.mlp_network_config
+            )
+
+        self.finite_difference_normal_eps: Optional[float] = None
+
+    def initialize_shape(self) -> None:
+        if self.cfg.shape_init is None and not self.cfg.force_shape_init:
+            return
+
+        # do not initialize shape if weights are provided
+        if self.cfg.weights is not None and not self.cfg.force_shape_init:
+            return
+
+        if self.cfg.sdf_bias != 0.0:
+            threestudio.warn(
+                "shape_init and sdf_bias are both specified, which may lead to unexpected results."
+            )
+
+        get_gt_sdf: Callable[[Float[Tensor, "N 3"]], Float[Tensor, "N 1"]]
+        assert isinstance(self.cfg.shape_init, str)
+        if self.cfg.shape_init == "ellipsoid":
+            assert (
+                isinstance(self.cfg.shape_init_params, Sized)
+                and len(self.cfg.shape_init_params) == 3
+            )
+            size = torch.as_tensor(self.cfg.shape_init_params).to(self.device)
+
+            def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
+                return ((points_rand / size) ** 2).sum(
+                    dim=-1, keepdim=True
+                ).sqrt() - 1.0  # pseudo signed distance of an ellipsoid
+
+            get_gt_sdf = func
+        elif self.cfg.shape_init == "sphere":
+            assert isinstance(self.cfg.shape_init_params, float)
+            radius = self.cfg.shape_init_params
+
+            def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
+                return (points_rand**2).sum(dim=-1, keepdim=True).sqrt() - radius
+
+            get_gt_sdf = func
+        elif self.cfg.shape_init.startswith("mesh:"):
+            assert isinstance(self.cfg.shape_init_params, float)
+            mesh_path = self.cfg.shape_init[5:]
+            if not os.path.exists(mesh_path):
+                raise ValueError(f"Mesh file {mesh_path} does not exist.")
+
+            import trimesh
+
+            scene = trimesh.load(mesh_path)
+            if isinstance(scene, trimesh.Trimesh):
+                mesh = scene
+            elif isinstance(scene, trimesh.scene.Scene):
+                mesh = trimesh.Trimesh()
+                for obj in scene.geometry.values():
+                    mesh = trimesh.util.concatenate([mesh, obj])
+            else:
+                raise ValueError(f"Unknown mesh type at {mesh_path}.")
+
+            # move to center
+            centroid = mesh.vertices.mean(0)
+            mesh.vertices = mesh.vertices - centroid
+
+            # align to up-z and front-x
+            dirs = ["+x", "+y", "+z", "-x", "-y", "-z"]
+            dir2vec = {
+                "+x": np.array([1, 0, 0]),
+                "+y": np.array([0, 1, 0]),
+                "+z": np.array([0, 0, 1]),
+                "-x": np.array([-1, 0, 0]),
+                "-y": np.array([0, -1, 0]),
+                "-z": np.array([0, 0, -1]),
+            }
+            if (
+                self.cfg.shape_init_mesh_up not in dirs
+                or self.cfg.shape_init_mesh_front not in dirs
+            ):
+                raise ValueError(
+                    f"shape_init_mesh_up and shape_init_mesh_front must be one of {dirs}."
+                )
+            if self.cfg.shape_init_mesh_up[1] == self.cfg.shape_init_mesh_front[1]:
+                raise ValueError(
+                    "shape_init_mesh_up and shape_init_mesh_front must be orthogonal."
+                )
+            z_, x_ = (
+                dir2vec[self.cfg.shape_init_mesh_up],
+                dir2vec[self.cfg.shape_init_mesh_front],
+            )
+            y_ = np.cross(z_, x_)
+            std2mesh = np.stack([x_, y_, z_], axis=0).T
+            mesh2std = np.linalg.inv(std2mesh)
+
+            # scaling
+            scale = np.abs(mesh.vertices).max()
+            mesh.vertices = mesh.vertices / scale * self.cfg.shape_init_params
+            mesh.vertices = np.dot(mesh2std, mesh.vertices.T).T
+
+            from pysdf import SDF
+
+            sdf = SDF(mesh.vertices, mesh.faces)
+
+            def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
+                # add a negative signed here
+                # as in pysdf the inside of the shape has positive signed distance
+                return torch.from_numpy(-sdf(points_rand.cpu().numpy())).to(
+                    points_rand
+                )[..., None]
+
+            get_gt_sdf = func
+
+        else:
+            raise ValueError(
+                f"Unknown shape initialization type: {self.cfg.shape_init}"
+            )
+
+        # Initialize SDF to a given shape when no weights are provided or force_shape_init is True
+        optim = torch.optim.Adam(self.parameters(), lr=1e-3)
+        from tqdm import tqdm
+
+        for _ in tqdm(
+            range(1000),
+            desc=f"Initializing SDF to a(n) {self.cfg.shape_init}:",
+            disable=get_rank() != 0,
+        ):
+            points_rand = (
+                torch.rand((10000, 3), dtype=torch.float32).to(self.device) * 2.0 - 1.0
+            )
+            sdf_gt = get_gt_sdf(points_rand)
+            sdf_pred = self.forward_sdf(points_rand)
+            loss = F.mse_loss(sdf_pred, sdf_gt)
+            optim.zero_grad()
+            loss.backward()
+            optim.step()
+
+        # explicit broadcast to ensure param consistency across ranks
+        for param in self.parameters():
+            broadcast(param, src=0)
+
+    def get_shifted_sdf(
+        self, points: Float[Tensor, "*N Di"], sdf: Float[Tensor, "*N 1"]
+    ) -> Float[Tensor, "*N 1"]:
+        sdf_bias: Union[float, Float[Tensor, "*N 1"]]
+        if self.cfg.sdf_bias == "ellipsoid":
+            assert (
+                isinstance(self.cfg.sdf_bias_params, Sized)
+                and len(self.cfg.sdf_bias_params) == 3
+            )
+            size = torch.as_tensor(self.cfg.sdf_bias_params).to(points)
+            sdf_bias = ((points / size) ** 2).sum(
+                dim=-1, keepdim=True
+            ).sqrt() - 1.0  # pseudo signed distance of an ellipsoid
+        elif self.cfg.sdf_bias == "sphere":
+            assert isinstance(self.cfg.sdf_bias_params, float)
+            radius = self.cfg.sdf_bias_params
+            sdf_bias = (points**2).sum(dim=-1, keepdim=True).sqrt() - radius
+        elif isinstance(self.cfg.sdf_bias, float):
+            sdf_bias = self.cfg.sdf_bias
+        else:
+            raise ValueError(f"Unknown sdf bias {self.cfg.sdf_bias}")
+        return sdf + sdf_bias
+
+    def forward(
+        self, points: Float[Tensor, "*N Di"], output_normal: bool = False
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        grad_enabled = torch.is_grad_enabled()
+
+        if output_normal and self.cfg.normal_type == "analytic":
+            torch.set_grad_enabled(True)
+            points.requires_grad_(True)
+
+        points_unscaled = points  # points in the original scale
+        points = contract_to_unisphere(
+            points, self.bbox, self.unbounded
+        )  # points normalized to (0, 1)
+
+        enc = self.encoding(points.view(-1, self.cfg.n_input_dims))
+        sdf = self.sdf_network(enc).view(*points.shape[:-1], 1)
+        sdf = self.get_shifted_sdf(points_unscaled, sdf)
+        output = {"sdf": sdf}
+
+        if self.cfg.n_feature_dims > 0:
+            features = self.feature_network(enc).view(
+                *points.shape[:-1], self.cfg.n_feature_dims
+            )
+            output.update({"features": features})
+
+        if output_normal:
+            if (
+                self.cfg.normal_type == "finite_difference"
+                or self.cfg.normal_type == "finite_difference_laplacian"
+            ):
+                assert self.finite_difference_normal_eps is not None
+                eps: float = self.finite_difference_normal_eps
+                if self.cfg.normal_type == "finite_difference_laplacian":
+                    offsets: Float[Tensor, "6 3"] = torch.as_tensor(
+                        [
+                            [eps, 0.0, 0.0],
+                            [-eps, 0.0, 0.0],
+                            [0.0, eps, 0.0],
+                            [0.0, -eps, 0.0],
+                            [0.0, 0.0, eps],
+                            [0.0, 0.0, -eps],
+                        ]
+                    ).to(points_unscaled)
+                    points_offset: Float[Tensor, "... 6 3"] = (
+                        points_unscaled[..., None, :] + offsets
+                    ).clamp(-self.cfg.radius, self.cfg.radius)
+                    sdf_offset: Float[Tensor, "... 6 1"] = self.forward_sdf(
+                        points_offset
+                    )
+                    sdf_grad = (
+                        0.5
+                        * (sdf_offset[..., 0::2, 0] - sdf_offset[..., 1::2, 0])
+                        / eps
+                    )
+                else:
+                    offsets: Float[Tensor, "3 3"] = torch.as_tensor(
+                        [[eps, 0.0, 0.0], [0.0, eps, 0.0], [0.0, 0.0, eps]]
+                    ).to(points_unscaled)
+                    points_offset: Float[Tensor, "... 3 3"] = (
+                        points_unscaled[..., None, :] + offsets
+                    ).clamp(-self.cfg.radius, self.cfg.radius)
+                    sdf_offset: Float[Tensor, "... 3 1"] = self.forward_sdf(
+                        points_offset
+                    )
+                    sdf_grad = (sdf_offset[..., 0::1, 0] - sdf) / eps
+                normal = F.normalize(sdf_grad, dim=-1)
+            elif self.cfg.normal_type == "pred":
+                normal = self.normal_network(enc).view(*points.shape[:-1], 3)
+                normal = F.normalize(normal, dim=-1)
+                sdf_grad = normal
+            elif self.cfg.normal_type == "analytic":
+                sdf_grad = -torch.autograd.grad(
+                    sdf,
+                    points_unscaled,
+                    grad_outputs=torch.ones_like(sdf),
+                    create_graph=True,
+                )[0]
+                normal = F.normalize(sdf_grad, dim=-1)
+                if not grad_enabled:
+                    sdf_grad = sdf_grad.detach()
+                    normal = normal.detach()
+            else:
+                raise AttributeError(f"Unknown normal type {self.cfg.normal_type}")
+            output.update(
+                {"normal": normal, "shading_normal": normal, "sdf_grad": sdf_grad}
+            )
+        return output
+
+    def forward_sdf(self, points: Float[Tensor, "*N Di"]) -> Float[Tensor, "*N 1"]:
+        points_unscaled = points
+        points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
+
+        sdf = self.sdf_network(
+            self.encoding(points.reshape(-1, self.cfg.n_input_dims))
+        ).reshape(*points.shape[:-1], 1)
+        sdf = self.get_shifted_sdf(points_unscaled, sdf)
+        return sdf
+
+    def forward_field(
+        self, points: Float[Tensor, "*N Di"]
+    ) -> Tuple[Float[Tensor, "*N 1"], Optional[Float[Tensor, "*N 3"]]]:
+        points_unscaled = points
+        points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
+        enc = self.encoding(points.reshape(-1, self.cfg.n_input_dims))
+        sdf = self.sdf_network(enc).reshape(*points.shape[:-1], 1)
+        sdf = self.get_shifted_sdf(points_unscaled, sdf)
+        deformation: Optional[Float[Tensor, "*N 3"]] = None
+        if self.cfg.isosurface_deformable_grid:
+            deformation = self.deformation_network(enc).reshape(*points.shape[:-1], 3)
+        return sdf, deformation
+
+    def forward_level(
+        self, field: Float[Tensor, "*N 1"], threshold: float
+    ) -> Float[Tensor, "*N 1"]:
+        return field - threshold
+
+    def export(self, points: Float[Tensor, "*N Di"], **kwargs) -> Dict[str, Any]:
+        out: Dict[str, Any] = {}
+        if self.cfg.n_feature_dims == 0:
+            return out
+        points_unscaled = points
+        points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
+        enc = self.encoding(points.reshape(-1, self.cfg.n_input_dims))
+        features = self.feature_network(enc).view(
+            *points.shape[:-1], self.cfg.n_feature_dims
+        )
+        out.update(
+            {
+                "features": features,
+            }
+        )
+        return out
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        if (
+            self.cfg.normal_type == "finite_difference"
+            or self.cfg.normal_type == "finite_difference_laplacian"
+        ):
+            if isinstance(self.cfg.finite_difference_normal_eps, float):
+                self.finite_difference_normal_eps = (
+                    self.cfg.finite_difference_normal_eps
+                )
+            elif self.cfg.finite_difference_normal_eps == "progressive":
+                # progressive finite difference eps from Neuralangelo
+                # https://arxiv.org/abs/2306.03092
+                hg_conf: Any = self.cfg.pos_encoding_config
+                assert (
+                    hg_conf.otype == "ProgressiveBandHashGrid"
+                ), "finite_difference_normal_eps=progressive only works with ProgressiveBandHashGrid"
+                current_level = min(
+                    hg_conf.start_level
+                    + max(global_step - hg_conf.start_step, 0) // hg_conf.update_steps,
+                    hg_conf.n_levels,
+                )
+                grid_res = hg_conf.base_resolution * hg_conf.per_level_scale ** (
+                    current_level - 1
+                )
+                grid_size = 2 * self.cfg.radius / grid_res
+                if grid_size != self.finite_difference_normal_eps:
+                    threestudio.info(
+                        f"Update finite_difference_normal_eps to {grid_size}"
+                    )
+                self.finite_difference_normal_eps = grid_size
+            else:
+                raise ValueError(
+                    f"Unknown finite_difference_normal_eps={self.cfg.finite_difference_normal_eps}"
+                )

threestudio/models/geometry/implicit_volume.py

@@ -0,0 +1,325 @@
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.geometry.base import (
+    BaseGeometry,
+    BaseImplicitGeometry,
+    contract_to_unisphere,
+)
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.ops import get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("implicit-volume")
+class ImplicitVolume(BaseImplicitGeometry):
+    @dataclass
+    class Config(BaseImplicitGeometry.Config):
+        n_input_dims: int = 3
+        n_feature_dims: int = 3
+        density_activation: Optional[str] = "softplus"
+        density_bias: Union[float, str] = "blob_magic3d"
+        density_blob_scale: float = 10.0
+        density_blob_std: float = 0.5
+        pos_encoding_config: dict = field(
+            default_factory=lambda: {
+                "otype": "HashGrid",
+                "n_levels": 16,
+                "n_features_per_level": 2,
+                "log2_hashmap_size": 19,
+                "base_resolution": 16,
+                "per_level_scale": 1.447269237440378,
+            }
+        )
+        mlp_network_config: dict = field(
+            default_factory=lambda: {
+                "otype": "VanillaMLP",
+                "activation": "ReLU",
+                "output_activation": "none",
+                "n_neurons": 64,
+                "n_hidden_layers": 1,
+            }
+        )
+        normal_type: Optional[
+            str
+        ] = "finite_difference"  # in ['pred', 'finite_difference', 'finite_difference_laplacian']
+        finite_difference_normal_eps: Union[
+            float, str
+        ] = 0.01  # in [float, "progressive"]
+
+        # automatically determine the threshold
+        isosurface_threshold: Union[float, str] = 25.0
+
+        # 4D Gaussian Annealing
+        anneal_density_blob_std_config: Optional[dict] = None
+
+    cfg: Config
+
+    def configure(self) -> None:
+        super().configure()
+        self.encoding = get_encoding(
+            self.cfg.n_input_dims, self.cfg.pos_encoding_config
+        )
+        self.density_network = get_mlp(
+            self.encoding.n_output_dims, 1, self.cfg.mlp_network_config
+        )
+        if self.cfg.n_feature_dims > 0:
+            self.feature_network = get_mlp(
+                self.encoding.n_output_dims,
+                self.cfg.n_feature_dims,
+                self.cfg.mlp_network_config,
+            )
+        if self.cfg.normal_type == "pred":
+            self.normal_network = get_mlp(
+                self.encoding.n_output_dims, 3, self.cfg.mlp_network_config
+            )
+
+        self.finite_difference_normal_eps: Optional[float] = None
+
+    def get_activated_density(
+        self, points: Float[Tensor, "*N Di"], density: Float[Tensor, "*N 1"]
+    ) -> Tuple[Float[Tensor, "*N 1"], Float[Tensor, "*N 1"]]:
+        density_bias: Union[float, Float[Tensor, "*N 1"]]
+        if self.cfg.density_bias == "blob_dreamfusion":
+            # pre-activation density bias
+            density_bias = (
+                self.cfg.density_blob_scale
+                * torch.exp(
+                    -0.5 * (points**2).sum(dim=-1) / self.cfg.density_blob_std**2
+                )[..., None]
+            )
+        elif self.cfg.density_bias == "blob_magic3d":
+            # pre-activation density bias
+            density_bias = (
+                self.cfg.density_blob_scale
+                * (
+                    1
+                    - torch.sqrt((points**2).sum(dim=-1)) / self.cfg.density_blob_std
+                )[..., None]
+            )
+        elif isinstance(self.cfg.density_bias, float):
+            density_bias = self.cfg.density_bias
+        else:
+            raise ValueError(f"Unknown density bias {self.cfg.density_bias}")
+        raw_density: Float[Tensor, "*N 1"] = density + density_bias
+        density = get_activation(self.cfg.density_activation)(raw_density)
+        return raw_density, density
+
+    def forward(
+        self, points: Float[Tensor, "*N Di"], output_normal: bool = False
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        grad_enabled = torch.is_grad_enabled()
+
+        if output_normal and self.cfg.normal_type == "analytic":
+            torch.set_grad_enabled(True)
+            points.requires_grad_(True)
+
+        points_unscaled = points  # points in the original scale
+        points = contract_to_unisphere(
+            points, self.bbox, self.unbounded
+        )  # points normalized to (0, 1)
+
+        enc = self.encoding(points.view(-1, self.cfg.n_input_dims))
+        density = self.density_network(enc).view(*points.shape[:-1], 1)
+        raw_density, density = self.get_activated_density(points_unscaled, density)
+
+        output = {
+            "density": density,
+        }
+
+        if self.cfg.n_feature_dims > 0:
+            features = self.feature_network(enc).view(
+                *points.shape[:-1], self.cfg.n_feature_dims
+            )
+            output.update({"features": features})
+
+        if output_normal:
+            if (
+                self.cfg.normal_type == "finite_difference"
+                or self.cfg.normal_type == "finite_difference_laplacian"
+            ):
+                # TODO: use raw density
+                assert self.finite_difference_normal_eps is not None
+                eps: float = self.finite_difference_normal_eps
+                if self.cfg.normal_type == "finite_difference_laplacian":
+                    offsets: Float[Tensor, "6 3"] = torch.as_tensor(
+                        [
+                            [eps, 0.0, 0.0],
+                            [-eps, 0.0, 0.0],
+                            [0.0, eps, 0.0],
+                            [0.0, -eps, 0.0],
+                            [0.0, 0.0, eps],
+                            [0.0, 0.0, -eps],
+                        ]
+                    ).to(points_unscaled)
+                    points_offset: Float[Tensor, "... 6 3"] = (
+                        points_unscaled[..., None, :] + offsets
+                    ).clamp(-self.cfg.radius, self.cfg.radius)
+                    density_offset: Float[Tensor, "... 6 1"] = self.forward_density(
+                        points_offset
+                    )
+                    normal = (
+                        -0.5
+                        * (density_offset[..., 0::2, 0] - density_offset[..., 1::2, 0])
+                        / eps
+                    )
+                else:
+                    offsets: Float[Tensor, "3 3"] = torch.as_tensor(
+                        [[eps, 0.0, 0.0], [0.0, eps, 0.0], [0.0, 0.0, eps]]
+                    ).to(points_unscaled)
+                    points_offset: Float[Tensor, "... 3 3"] = (
+                        points_unscaled[..., None, :] + offsets
+                    ).clamp(-self.cfg.radius, self.cfg.radius)
+                    density_offset: Float[Tensor, "... 3 1"] = self.forward_density(
+                        points_offset
+                    )
+                    normal = -(density_offset[..., 0::1, 0] - density) / eps
+                normal = F.normalize(normal, dim=-1)
+            elif self.cfg.normal_type == "pred":
+                normal = self.normal_network(enc).view(*points.shape[:-1], 3)
+                normal = F.normalize(normal, dim=-1)
+            elif self.cfg.normal_type == "analytic":
+                normal = -torch.autograd.grad(
+                    density,
+                    points_unscaled,
+                    grad_outputs=torch.ones_like(density),
+                    create_graph=True,
+                )[0]
+                normal = F.normalize(normal, dim=-1)
+                if not grad_enabled:
+                    normal = normal.detach()
+            else:
+                raise AttributeError(f"Unknown normal type {self.cfg.normal_type}")
+            output.update({"normal": normal, "shading_normal": normal})
+
+        torch.set_grad_enabled(grad_enabled)
+        return output
+
+    def forward_density(self, points: Float[Tensor, "*N Di"]) -> Float[Tensor, "*N 1"]:
+        points_unscaled = points
+        points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
+
+        density = self.density_network(
+            self.encoding(points.reshape(-1, self.cfg.n_input_dims))
+        ).reshape(*points.shape[:-1], 1)
+
+        _, density = self.get_activated_density(points_unscaled, density)
+        return density
+
+    def forward_field(
+        self, points: Float[Tensor, "*N Di"]
+    ) -> Tuple[Float[Tensor, "*N 1"], Optional[Float[Tensor, "*N 3"]]]:
+        if self.cfg.isosurface_deformable_grid:
+            threestudio.warn(
+                f"{self.__class__.__name__} does not support isosurface_deformable_grid. Ignoring."
+            )
+        density = self.forward_density(points)
+        return density, None
+
+    def forward_level(
+        self, field: Float[Tensor, "*N 1"], threshold: float
+    ) -> Float[Tensor, "*N 1"]:
+        return -(field - threshold)
+
+    def export(self, points: Float[Tensor, "*N Di"], **kwargs) -> Dict[str, Any]:
+        out: Dict[str, Any] = {}
+        if self.cfg.n_feature_dims == 0:
+            return out
+        points_unscaled = points
+        points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
+        enc = self.encoding(points.reshape(-1, self.cfg.n_input_dims))
+        features = self.feature_network(enc).view(
+            *points.shape[:-1], self.cfg.n_feature_dims
+        )
+        out.update(
+            {
+                "features": features,
+            }
+        )
+        return out
+
+    @staticmethod
+    @torch.no_grad()
+    def create_from(
+        other: BaseGeometry,
+        cfg: Optional[Union[dict, DictConfig]] = None,
+        copy_net: bool = True,
+        **kwargs,
+    ) -> "ImplicitVolume":
+        if isinstance(other, ImplicitVolume):
+            instance = ImplicitVolume(cfg, **kwargs)
+            instance.encoding.load_state_dict(other.encoding.state_dict())
+            instance.density_network.load_state_dict(other.density_network.state_dict())
+            if copy_net:
+                if (
+                    instance.cfg.n_feature_dims > 0
+                    and other.cfg.n_feature_dims == instance.cfg.n_feature_dims
+                ):
+                    instance.feature_network.load_state_dict(
+                        other.feature_network.state_dict()
+                    )
+                if (
+                    instance.cfg.normal_type == "pred"
+                    and other.cfg.normal_type == "pred"
+                ):
+                    instance.normal_network.load_state_dict(
+                        other.normal_network.state_dict()
+                    )
+            return instance
+        else:
+            raise TypeError(
+                f"Cannot create {ImplicitVolume.__name__} from {other.__class__.__name__}"
+            )
+
+    # FIXME: use progressive normal eps
+    def update_step(
+        self, epoch: int, global_step: int, on_load_weights: bool = False
+    ) -> None:
+        if self.cfg.anneal_density_blob_std_config is not None:
+            min_step = self.cfg.anneal_density_blob_std_config.min_anneal_step
+            max_step = self.cfg.anneal_density_blob_std_config.max_anneal_step
+            if global_step >= min_step and global_step <= max_step:
+                end_val = self.cfg.anneal_density_blob_std_config.end_val
+                start_val = self.cfg.anneal_density_blob_std_config.start_val
+                self.density_blob_std = start_val + (global_step - min_step) * (
+                    end_val - start_val
+                ) / (max_step - min_step)
+
+        if (
+            self.cfg.normal_type == "finite_difference"
+            or self.cfg.normal_type == "finite_difference_laplacian"
+        ):
+            if isinstance(self.cfg.finite_difference_normal_eps, float):
+                self.finite_difference_normal_eps = (
+                    self.cfg.finite_difference_normal_eps
+                )
+            elif self.cfg.finite_difference_normal_eps == "progressive":
+                # progressive finite difference eps from Neuralangelo
+                # https://arxiv.org/abs/2306.03092
+                hg_conf: Any = self.cfg.pos_encoding_config
+                assert (
+                    hg_conf.otype == "ProgressiveBandHashGrid"
+                ), "finite_difference_normal_eps=progressive only works with ProgressiveBandHashGrid"
+                current_level = min(
+                    hg_conf.start_level
+                    + max(global_step - hg_conf.start_step, 0) // hg_conf.update_steps,
+                    hg_conf.n_levels,
+                )
+                grid_res = hg_conf.base_resolution * hg_conf.per_level_scale ** (
+                    current_level - 1
+                )
+                grid_size = 2 * self.cfg.radius / grid_res
+                if grid_size != self.finite_difference_normal_eps:
+                    threestudio.info(
+                        f"Update finite_difference_normal_eps to {grid_size}"
+                    )
+                self.finite_difference_normal_eps = grid_size
+            else:
+                raise ValueError(
+                    f"Unknown finite_difference_normal_eps={self.cfg.finite_difference_normal_eps}"
+                )

threestudio/models/geometry/tetrahedra_sdf_grid.py

@@ -0,0 +1,369 @@
+import os
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.geometry.base import (
+    BaseExplicitGeometry,
+    BaseGeometry,
+    contract_to_unisphere,
+)
+from threestudio.models.geometry.implicit_sdf import ImplicitSDF
+from threestudio.models.geometry.implicit_volume import ImplicitVolume
+from threestudio.models.isosurface import MarchingTetrahedraHelper
+from threestudio.models.mesh import Mesh
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.misc import broadcast
+from threestudio.utils.ops import scale_tensor
+from threestudio.utils.typing import *
+
+
+@threestudio.register("tetrahedra-sdf-grid")
+class TetrahedraSDFGrid(BaseExplicitGeometry):
+    @dataclass
+    class Config(BaseExplicitGeometry.Config):
+        isosurface_resolution: int = 128
+        isosurface_deformable_grid: bool = True
+        isosurface_remove_outliers: bool = False
+        isosurface_outlier_n_faces_threshold: Union[int, float] = 0.01
+
+        n_input_dims: int = 3
+        n_feature_dims: int = 3
+        pos_encoding_config: dict = field(
+            default_factory=lambda: {
+                "otype": "HashGrid",
+                "n_levels": 16,
+                "n_features_per_level": 2,
+                "log2_hashmap_size": 19,
+                "base_resolution": 16,
+                "per_level_scale": 1.447269237440378,
+            }
+        )
+        mlp_network_config: dict = field(
+            default_factory=lambda: {
+                "otype": "VanillaMLP",
+                "activation": "ReLU",
+                "output_activation": "none",
+                "n_neurons": 64,
+                "n_hidden_layers": 1,
+            }
+        )
+        shape_init: Optional[str] = None
+        shape_init_params: Optional[Any] = None
+        shape_init_mesh_up: str = "+z"
+        shape_init_mesh_front: str = "+x"
+        force_shape_init: bool = False
+        geometry_only: bool = False
+        fix_geometry: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        super().configure()
+
+        # this should be saved to state_dict, register as buffer
+        self.isosurface_bbox: Float[Tensor, "2 3"]
+        self.register_buffer("isosurface_bbox", self.bbox.clone())
+
+        self.isosurface_helper = MarchingTetrahedraHelper(
+            self.cfg.isosurface_resolution,
+            f"load/tets/{self.cfg.isosurface_resolution}_tets.npz",
+        )
+
+        self.sdf: Float[Tensor, "Nv 1"]
+        self.deformation: Optional[Float[Tensor, "Nv 3"]]
+
+        if not self.cfg.fix_geometry:
+            self.register_parameter(
+                "sdf",
+                nn.Parameter(
+                    torch.zeros(
+                        (self.isosurface_helper.grid_vertices.shape[0], 1),
+                        dtype=torch.float32,
+                    )
+                ),
+            )
+            if self.cfg.isosurface_deformable_grid:
+                self.register_parameter(
+                    "deformation",
+                    nn.Parameter(
+                        torch.zeros_like(self.isosurface_helper.grid_vertices)
+                    ),
+                )
+            else:
+                self.deformation = None
+        else:
+            self.register_buffer(
+                "sdf",
+                torch.zeros(
+                    (self.isosurface_helper.grid_vertices.shape[0], 1),
+                    dtype=torch.float32,
+                ),
+            )
+            if self.cfg.isosurface_deformable_grid:
+                self.register_buffer(
+                    "deformation",
+                    torch.zeros_like(self.isosurface_helper.grid_vertices),
+                )
+            else:
+                self.deformation = None
+
+        if not self.cfg.geometry_only:
+            self.encoding = get_encoding(
+                self.cfg.n_input_dims, self.cfg.pos_encoding_config
+            )
+            self.feature_network = get_mlp(
+                self.encoding.n_output_dims,
+                self.cfg.n_feature_dims,
+                self.cfg.mlp_network_config,
+            )
+
+        self.mesh: Optional[Mesh] = None
+
+    def initialize_shape(self) -> None:
+        if self.cfg.shape_init is None and not self.cfg.force_shape_init:
+            return
+
+        # do not initialize shape if weights are provided
+        if self.cfg.weights is not None and not self.cfg.force_shape_init:
+            return
+
+        get_gt_sdf: Callable[[Float[Tensor, "N 3"]], Float[Tensor, "N 1"]]
+        assert isinstance(self.cfg.shape_init, str)
+        if self.cfg.shape_init == "ellipsoid":
+            assert (
+                isinstance(self.cfg.shape_init_params, Sized)
+                and len(self.cfg.shape_init_params) == 3
+            )
+            size = torch.as_tensor(self.cfg.shape_init_params).to(self.device)
+
+            def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
+                return ((points_rand / size) ** 2).sum(
+                    dim=-1, keepdim=True
+                ).sqrt() - 1.0  # pseudo signed distance of an ellipsoid
+
+            get_gt_sdf = func
+        elif self.cfg.shape_init == "sphere":
+            assert isinstance(self.cfg.shape_init_params, float)
+            radius = self.cfg.shape_init_params
+
+            def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
+                return (points_rand**2).sum(dim=-1, keepdim=True).sqrt() - radius
+
+            get_gt_sdf = func
+        elif self.cfg.shape_init.startswith("mesh:"):
+            assert isinstance(self.cfg.shape_init_params, float)
+            mesh_path = self.cfg.shape_init[5:]
+            if not os.path.exists(mesh_path):
+                raise ValueError(f"Mesh file {mesh_path} does not exist.")
+
+            import trimesh
+
+            mesh = trimesh.load(mesh_path)
+
+            # move to center
+            centroid = mesh.vertices.mean(0)
+            mesh.vertices = mesh.vertices - centroid
+
+            # align to up-z and front-x
+            dirs = ["+x", "+y", "+z", "-x", "-y", "-z"]
+            dir2vec = {
+                "+x": np.array([1, 0, 0]),
+                "+y": np.array([0, 1, 0]),
+                "+z": np.array([0, 0, 1]),
+                "-x": np.array([-1, 0, 0]),
+                "-y": np.array([0, -1, 0]),
+                "-z": np.array([0, 0, -1]),
+            }
+            if (
+                self.cfg.shape_init_mesh_up not in dirs
+                or self.cfg.shape_init_mesh_front not in dirs
+            ):
+                raise ValueError(
+                    f"shape_init_mesh_up and shape_init_mesh_front must be one of {dirs}."
+                )
+            if self.cfg.shape_init_mesh_up[1] == self.cfg.shape_init_mesh_front[1]:
+                raise ValueError(
+                    "shape_init_mesh_up and shape_init_mesh_front must be orthogonal."
+                )
+            z_, x_ = (
+                dir2vec[self.cfg.shape_init_mesh_up],
+                dir2vec[self.cfg.shape_init_mesh_front],
+            )
+            y_ = np.cross(z_, x_)
+            std2mesh = np.stack([x_, y_, z_], axis=0).T
+            mesh2std = np.linalg.inv(std2mesh)
+
+            # scaling
+            scale = np.abs(mesh.vertices).max()
+            mesh.vertices = mesh.vertices / scale * self.cfg.shape_init_params
+            mesh.vertices = np.dot(mesh2std, mesh.vertices.T).T
+
+            from pysdf import SDF
+
+            sdf = SDF(mesh.vertices, mesh.faces)
+
+            def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
+                # add a negative signed here
+                # as in pysdf the inside of the shape has positive signed distance
+                return torch.from_numpy(-sdf(points_rand.cpu().numpy())).to(
+                    points_rand
+                )[..., None]
+
+            get_gt_sdf = func
+
+        else:
+            raise ValueError(
+                f"Unknown shape initialization type: {self.cfg.shape_init}"
+            )
+
+        sdf_gt = get_gt_sdf(
+            scale_tensor(
+                self.isosurface_helper.grid_vertices,
+                self.isosurface_helper.points_range,
+                self.isosurface_bbox,
+            )
+        )
+        self.sdf.data = sdf_gt
+
+        # explicit broadcast to ensure param consistency across ranks
+        for param in self.parameters():
+            broadcast(param, src=0)
+
+    def isosurface(self) -> Mesh:
+        # return cached mesh if fix_geometry is True to save computation
+        if self.cfg.fix_geometry and self.mesh is not None:
+            return self.mesh
+        mesh = self.isosurface_helper(self.sdf, self.deformation)
+        mesh.v_pos = scale_tensor(
+            mesh.v_pos, self.isosurface_helper.points_range, self.isosurface_bbox
+        )
+        if self.cfg.isosurface_remove_outliers:
+            mesh = mesh.remove_outlier(self.cfg.isosurface_outlier_n_faces_threshold)
+        self.mesh = mesh
+        return mesh
+
+    def forward(
+        self, points: Float[Tensor, "*N Di"], output_normal: bool = False
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        if self.cfg.geometry_only:
+            return {}
+        assert (
+            output_normal == False
+        ), f"Normal output is not supported for {self.__class__.__name__}"
+        points_unscaled = points  # points in the original scale
+        points = contract_to_unisphere(points, self.bbox)  # points normalized to (0, 1)
+        enc = self.encoding(points.view(-1, self.cfg.n_input_dims))
+        features = self.feature_network(enc).view(
+            *points.shape[:-1], self.cfg.n_feature_dims
+        )
+        return {"features": features}
+
+    @staticmethod
+    @torch.no_grad()
+    def create_from(
+        other: BaseGeometry,
+        cfg: Optional[Union[dict, DictConfig]] = None,
+        copy_net: bool = True,
+        **kwargs,
+    ) -> "TetrahedraSDFGrid":
+        if isinstance(other, TetrahedraSDFGrid):
+            instance = TetrahedraSDFGrid(cfg, **kwargs)
+            assert instance.cfg.isosurface_resolution == other.cfg.isosurface_resolution
+            instance.isosurface_bbox = other.isosurface_bbox.clone()
+            instance.sdf.data = other.sdf.data.clone()
+            if (
+                instance.cfg.isosurface_deformable_grid
+                and other.cfg.isosurface_deformable_grid
+            ):
+                assert (
+                    instance.deformation is not None and other.deformation is not None
+                )
+                instance.deformation.data = other.deformation.data.clone()
+            if (
+                not instance.cfg.geometry_only
+                and not other.cfg.geometry_only
+                and copy_net
+            ):
+                instance.encoding.load_state_dict(other.encoding.state_dict())
+                instance.feature_network.load_state_dict(
+                    other.feature_network.state_dict()
+                )
+            return instance
+        elif isinstance(other, ImplicitVolume):
+            instance = TetrahedraSDFGrid(cfg, **kwargs)
+            if other.cfg.isosurface_method != "mt":
+                other.cfg.isosurface_method = "mt"
+                threestudio.warn(
+                    f"Override isosurface_method of the source geometry to 'mt'"
+                )
+            if other.cfg.isosurface_resolution != instance.cfg.isosurface_resolution:
+                other.cfg.isosurface_resolution = instance.cfg.isosurface_resolution
+                threestudio.warn(
+                    f"Override isosurface_resolution of the source geometry to {instance.cfg.isosurface_resolution}"
+                )
+            mesh = other.isosurface()
+            instance.isosurface_bbox = mesh.extras["bbox"]
+            instance.sdf.data = (
+                mesh.extras["grid_level"].to(instance.sdf.data).clamp(-1, 1)
+            )
+            if not instance.cfg.geometry_only and copy_net:
+                instance.encoding.load_state_dict(other.encoding.state_dict())
+                instance.feature_network.load_state_dict(
+                    other.feature_network.state_dict()
+                )
+            return instance
+        elif isinstance(other, ImplicitSDF):
+            instance = TetrahedraSDFGrid(cfg, **kwargs)
+            if other.cfg.isosurface_method != "mt":
+                other.cfg.isosurface_method = "mt"
+                threestudio.warn(
+                    f"Override isosurface_method of the source geometry to 'mt'"
+                )
+            if other.cfg.isosurface_resolution != instance.cfg.isosurface_resolution:
+                other.cfg.isosurface_resolution = instance.cfg.isosurface_resolution
+                threestudio.warn(
+                    f"Override isosurface_resolution of the source geometry to {instance.cfg.isosurface_resolution}"
+                )
+            mesh = other.isosurface()
+            instance.isosurface_bbox = mesh.extras["bbox"]
+            instance.sdf.data = mesh.extras["grid_level"].to(instance.sdf.data)
+            if (
+                instance.cfg.isosurface_deformable_grid
+                and other.cfg.isosurface_deformable_grid
+            ):
+                assert instance.deformation is not None
+                instance.deformation.data = mesh.extras["grid_deformation"].to(
+                    instance.deformation.data
+                )
+            if not instance.cfg.geometry_only and copy_net:
+                instance.encoding.load_state_dict(other.encoding.state_dict())
+                instance.feature_network.load_state_dict(
+                    other.feature_network.state_dict()
+                )
+            return instance
+        else:
+            raise TypeError(
+                f"Cannot create {TetrahedraSDFGrid.__name__} from {other.__class__.__name__}"
+            )
+
+    def export(self, points: Float[Tensor, "*N Di"], **kwargs) -> Dict[str, Any]:
+        out: Dict[str, Any] = {}
+        if self.cfg.geometry_only or self.cfg.n_feature_dims == 0:
+            return out
+        points_unscaled = points
+        points = contract_to_unisphere(points_unscaled, self.bbox)
+        enc = self.encoding(points.reshape(-1, self.cfg.n_input_dims))
+        features = self.feature_network(enc).view(
+            *points.shape[:-1], self.cfg.n_feature_dims
+        )
+        out.update(
+            {
+                "features": features,
+            }
+        )
+        return out

threestudio/models/geometry/volume_grid.py

@@ -0,0 +1,190 @@
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.geometry.base import BaseImplicitGeometry, contract_to_unisphere
+from threestudio.utils.ops import get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("volume-grid")
+class VolumeGrid(BaseImplicitGeometry):
+    @dataclass
+    class Config(BaseImplicitGeometry.Config):
+        grid_size: Tuple[int, int, int] = field(default_factory=lambda: (100, 100, 100))
+        n_feature_dims: int = 3
+        density_activation: Optional[str] = "softplus"
+        density_bias: Union[float, str] = "blob"
+        density_blob_scale: float = 5.0
+        density_blob_std: float = 0.5
+        normal_type: Optional[
+            str
+        ] = "finite_difference"  # in ['pred', 'finite_difference', 'finite_difference_laplacian']
+
+        # automatically determine the threshold
+        isosurface_threshold: Union[float, str] = "auto"
+
+    cfg: Config
+
+    def configure(self) -> None:
+        super().configure()
+        self.grid_size = self.cfg.grid_size
+
+        self.grid = nn.Parameter(
+            torch.zeros(1, self.cfg.n_feature_dims + 1, *self.grid_size)
+        )
+        if self.cfg.density_bias == "blob":
+            self.register_buffer("density_scale", torch.tensor(0.0))
+        else:
+            self.density_scale = nn.Parameter(torch.tensor(0.0))
+
+        if self.cfg.normal_type == "pred":
+            self.normal_grid = nn.Parameter(torch.zeros(1, 3, *self.grid_size))
+
+    def get_density_bias(self, points: Float[Tensor, "*N Di"]):
+        if self.cfg.density_bias == "blob":
+            # density_bias: Float[Tensor, "*N 1"] = self.cfg.density_blob_scale * torch.exp(-0.5 * (points ** 2).sum(dim=-1) / self.cfg.density_blob_std ** 2)[...,None]
+            density_bias: Float[Tensor, "*N 1"] = (
+                self.cfg.density_blob_scale
+                * (
+                    1
+                    - torch.sqrt((points.detach() ** 2).sum(dim=-1))
+                    / self.cfg.density_blob_std
+                )[..., None]
+            )
+            return density_bias
+        elif isinstance(self.cfg.density_bias, float):
+            return self.cfg.density_bias
+        else:
+            raise AttributeError(f"Unknown density bias {self.cfg.density_bias}")
+
+    def get_trilinear_feature(
+        self, points: Float[Tensor, "*N Di"], grid: Float[Tensor, "1 Df G1 G2 G3"]
+    ) -> Float[Tensor, "*N Df"]:
+        points_shape = points.shape[:-1]
+        df = grid.shape[1]
+        di = points.shape[-1]
+        out = F.grid_sample(
+            grid, points.view(1, 1, 1, -1, di), align_corners=False, mode="bilinear"
+        )
+        out = out.reshape(df, -1).T.reshape(*points_shape, df)
+        return out
+
+    def forward(
+        self, points: Float[Tensor, "*N Di"], output_normal: bool = False
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        points_unscaled = points  # points in the original scale
+        points = contract_to_unisphere(
+            points, self.bbox, self.unbounded
+        )  # points normalized to (0, 1)
+        points = points * 2 - 1  # convert to [-1, 1] for grid sample
+
+        out = self.get_trilinear_feature(points, self.grid)
+        density, features = out[..., 0:1], out[..., 1:]
+        density = density * torch.exp(self.density_scale)  # exp scaling in DreamFusion
+
+        # breakpoint()
+        density = get_activation(self.cfg.density_activation)(
+            density + self.get_density_bias(points_unscaled)
+        )
+
+        output = {
+            "density": density,
+            "features": features,
+        }
+
+        if output_normal:
+            if (
+                self.cfg.normal_type == "finite_difference"
+                or self.cfg.normal_type == "finite_difference_laplacian"
+            ):
+                eps = 1.0e-3
+                if self.cfg.normal_type == "finite_difference_laplacian":
+                    offsets: Float[Tensor, "6 3"] = torch.as_tensor(
+                        [
+                            [eps, 0.0, 0.0],
+                            [-eps, 0.0, 0.0],
+                            [0.0, eps, 0.0],
+                            [0.0, -eps, 0.0],
+                            [0.0, 0.0, eps],
+                            [0.0, 0.0, -eps],
+                        ]
+                    ).to(points_unscaled)
+                    points_offset: Float[Tensor, "... 6 3"] = (
+                        points_unscaled[..., None, :] + offsets
+                    ).clamp(-self.cfg.radius, self.cfg.radius)
+                    density_offset: Float[Tensor, "... 6 1"] = self.forward_density(
+                        points_offset
+                    )
+                    normal = (
+                        -0.5
+                        * (density_offset[..., 0::2, 0] - density_offset[..., 1::2, 0])
+                        / eps
+                    )
+                else:
+                    offsets: Float[Tensor, "3 3"] = torch.as_tensor(
+                        [[eps, 0.0, 0.0], [0.0, eps, 0.0], [0.0, 0.0, eps]]
+                    ).to(points_unscaled)
+                    points_offset: Float[Tensor, "... 3 3"] = (
+                        points_unscaled[..., None, :] + offsets
+                    ).clamp(-self.cfg.radius, self.cfg.radius)
+                    density_offset: Float[Tensor, "... 3 1"] = self.forward_density(
+                        points_offset
+                    )
+                    normal = -(density_offset[..., 0::1, 0] - density) / eps
+                normal = F.normalize(normal, dim=-1)
+            elif self.cfg.normal_type == "pred":
+                normal = self.get_trilinear_feature(points, self.normal_grid)
+                normal = F.normalize(normal, dim=-1)
+            else:
+                raise AttributeError(f"Unknown normal type {self.cfg.normal_type}")
+            output.update({"normal": normal, "shading_normal": normal})
+        return output
+
+    def forward_density(self, points: Float[Tensor, "*N Di"]) -> Float[Tensor, "*N 1"]:
+        points_unscaled = points
+        points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
+        points = points * 2 - 1  # convert to [-1, 1] for grid sample
+
+        out = self.get_trilinear_feature(points, self.grid)
+        density = out[..., 0:1]
+        density = density * torch.exp(self.density_scale)
+
+        density = get_activation(self.cfg.density_activation)(
+            density + self.get_density_bias(points_unscaled)
+        )
+        return density
+
+    def forward_field(
+        self, points: Float[Tensor, "*N Di"]
+    ) -> Tuple[Float[Tensor, "*N 1"], Optional[Float[Tensor, "*N 3"]]]:
+        if self.cfg.isosurface_deformable_grid:
+            threestudio.warn(
+                f"{self.__class__.__name__} does not support isosurface_deformable_grid. Ignoring."
+            )
+        density = self.forward_density(points)
+        return density, None
+
+    def forward_level(
+        self, field: Float[Tensor, "*N 1"], threshold: float
+    ) -> Float[Tensor, "*N 1"]:
+        return -(field - threshold)
+
+    def export(self, points: Float[Tensor, "*N Di"], **kwargs) -> Dict[str, Any]:
+        out: Dict[str, Any] = {}
+        if self.cfg.n_feature_dims == 0:
+            return out
+        points_unscaled = points
+        points = contract_to_unisphere(points, self.bbox, self.unbounded)
+        points = points * 2 - 1  # convert to [-1, 1] for grid sample
+        features = self.get_trilinear_feature(points, self.grid)[..., 1:]
+        out.update(
+            {
+                "features": features,
+            }
+        )
+        return out

threestudio/models/guidance/__init__.py

@@ -0,0 +1,13 @@
+from . import (
+    controlnet_guidance,
+    controlnet_reg_guidance,
+    deep_floyd_guidance,
+    stable_diffusion_guidance,
+    stable_diffusion_unified_guidance,
+    stable_diffusion_vsd_guidance,
+    stable_diffusion_bsd_guidance,
+    stable_zero123_guidance,
+    zero123_guidance,
+    zero123_unified_guidance,
+    clip_guidance,
+)

threestudio/models/guidance/clip_guidance.py

@@ -0,0 +1,84 @@
+from dataclasses import dataclass
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as T
+import clip
+
+import threestudio
+from threestudio.utils.base import BaseObject
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.typing import *
+
+
+@threestudio.register("clip-guidance")
+class CLIPGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        pretrained_model_name_or_path: str = "ViT-B/16"
+        view_dependent_prompting: bool = True
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading CLIP ...")
+        self.clip_model, self.clip_preprocess = clip.load(
+            self.cfg.pretrained_model_name_or_path,
+            device=self.device,
+            jit=False,
+            download_root=self.cfg.cache_dir
+        )
+
+        self.aug = T.Compose([
+            T.Resize((224, 224)),
+            T.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+        ])
+
+        threestudio.info(f"Loaded CLIP!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def get_embedding(self, input_value, is_text=True):
+        if is_text:
+            value = clip.tokenize(input_value).to(self.device)
+            z = self.clip_model.encode_text(value)
+        else:
+            input_value = self.aug(input_value)
+            z = self.clip_model.encode_image(input_value)
+
+        return z / z.norm(dim=-1, keepdim=True)
+
+    def get_loss(self, image_z, clip_z, loss_type='similarity_score', use_mean=True):
+        if loss_type == 'similarity_score':
+            loss = -((image_z * clip_z).sum(-1))
+        elif loss_type == 'spherical_dist':
+            image_z, clip_z = F.normalize(image_z, dim=-1), F.normalize(clip_z, dim=-1)
+            loss = ((image_z - clip_z).norm(dim=-1).div(2).arcsin().pow(2).mul(2))
+        else:
+            raise NotImplementedError
+
+        return loss.mean() if use_mean else loss
+
+    def __call__(
+        self,
+        pred_rgb: Float[Tensor, "B H W C"],
+        gt_rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        embedding_type: str = 'both',
+        loss_type: Optional[str] = 'similarity_score',
+        **kwargs,
+    ):
+        clip_text_loss, clip_img_loss = 0, 0
+
+        if embedding_type in ('both', 'text'):
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            ).chunk(2)[0]
+            clip_text_loss = self.get_loss(self.get_embedding(pred_rgb, is_text=False), text_embeddings, loss_type=loss_type)
+
+        if embedding_type in ('both', 'img'):
+            clip_img_loss = self.get_loss(self.get_embedding(pred_rgb, is_text=False), self.get_embedding(gt_rgb, is_text=False), loss_type=loss_type)
+
+        return clip_text_loss + clip_img_loss

threestudio/models/guidance/controlnet_guidance.py

@@ -0,0 +1,517 @@
+import os
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from controlnet_aux import CannyDetector, NormalBaeDetector
+from diffusers import ControlNetModel, DDIMScheduler, StableDiffusionControlNetPipeline
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.perceptual import PerceptualLoss
+from threestudio.utils.typing import *
+
+
+@threestudio.register("stable-diffusion-controlnet-guidance")
+class ControlNetGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        pretrained_model_name_or_path: str = "SG161222/Realistic_Vision_V2.0"
+        ddim_scheduler_name_or_path: str = "runwayml/stable-diffusion-v1-5"
+        control_type: str = "normal"  # normal/canny
+
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = False
+        guidance_scale: float = 7.5
+        condition_scale: float = 1.5
+        grad_clip: Optional[Any] = None
+        half_precision_weights: bool = True
+
+        fixed_size: int = -1
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+        diffusion_steps: int = 20
+
+        use_sds: bool = False
+
+        use_du: bool = False
+        per_du_step: int = 10
+        start_du_step: int = 1000
+        cache_du: bool = False
+
+        # Canny threshold
+        canny_lower_bound: int = 50
+        canny_upper_bound: int = 100
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading ControlNet ...")
+
+        controlnet_name_or_path: str
+        if self.cfg.control_type in ("normal", "input_normal"):
+            controlnet_name_or_path = "lllyasviel/control_v11p_sd15_normalbae"
+        elif self.cfg.control_type == "canny":
+            controlnet_name_or_path = "lllyasviel/control_v11p_sd15_canny"
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        pipe_kwargs = {
+            "safety_checker": None,
+            "feature_extractor": None,
+            "requires_safety_checker": False,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+        }
+
+        controlnet = ControlNetModel.from_pretrained(
+            controlnet_name_or_path,
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+        )
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path, controlnet=controlnet, **pipe_kwargs
+        ).to(self.device)
+        self.scheduler = DDIMScheduler.from_pretrained(
+            self.cfg.ddim_scheduler_name_or_path,
+            subfolder="scheduler",
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+        )
+        self.scheduler.set_timesteps(self.cfg.diffusion_steps)
+
+        if self.cfg.enable_memory_efficient_attention:
+            if parse_version(torch.__version__) >= parse_version("2"):
+                threestudio.info(
+                    "PyTorch2.0 uses memory efficient attention by default."
+                )
+            elif not is_xformers_available():
+                threestudio.warn(
+                    "xformers is not available, memory efficient attention is not enabled."
+                )
+            else:
+                self.pipe.enable_xformers_memory_efficient_attention()
+
+        if self.cfg.enable_sequential_cpu_offload:
+            self.pipe.enable_sequential_cpu_offload()
+
+        if self.cfg.enable_attention_slicing:
+            self.pipe.enable_attention_slicing(1)
+
+        if self.cfg.enable_channels_last_format:
+            self.pipe.unet.to(memory_format=torch.channels_last)
+
+        # Create model
+        self.vae = self.pipe.vae.eval()
+        self.unet = self.pipe.unet.eval()
+        self.controlnet = self.pipe.controlnet.eval()
+
+        if self.cfg.control_type == "normal":
+            self.preprocessor = NormalBaeDetector.from_pretrained(
+                "lllyasviel/Annotators"
+            )
+            self.preprocessor.model.to(self.device)
+        elif self.cfg.control_type == "canny":
+            self.preprocessor = CannyDetector()
+
+        for p in self.vae.parameters():
+            p.requires_grad_(False)
+        for p in self.unet.parameters():
+            p.requires_grad_(False)
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.grad_clip_val: Optional[float] = None
+
+        if self.cfg.use_du:
+            if self.cfg.cache_du:
+                self.edit_frames = {}
+            self.perceptual_loss = PerceptualLoss().eval().to(self.device)
+
+        threestudio.info(f"Loaded ControlNet!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_controlnet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        image_cond: Float[Tensor, "..."],
+        condition_scale: float,
+        encoder_hidden_states: Float[Tensor, "..."],
+    ) -> Float[Tensor, "..."]:
+        return self.controlnet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+            controlnet_cond=image_cond.to(self.weights_dtype),
+            conditioning_scale=condition_scale,
+            return_dict=False,
+        )
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_control_unet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+        cross_attention_kwargs,
+        down_block_additional_residuals,
+        mid_block_additional_residual,
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        return self.unet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+        ).sample.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 H W"]
+    ) -> Float[Tensor, "B 4 DH DW"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_cond_images(
+        self, imgs: Float[Tensor, "B 3 H W"]
+    ) -> Float[Tensor, "B 4 DH DW"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.mode()
+        uncond_image_latents = torch.zeros_like(latents)
+        latents = torch.cat([latents, latents, uncond_image_latents], dim=0)
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self, latents: Float[Tensor, "B 4 DH DW"]
+    ) -> Float[Tensor, "B 3 H W"]:
+        input_dtype = latents.dtype
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents.to(self.weights_dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    def edit_latents(
+        self,
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        latents: Float[Tensor, "B 4 DH DW"],
+        image_cond: Float[Tensor, "B 3 H W"],
+        t: Int[Tensor, "B"],
+        mask = None
+    ) -> Float[Tensor, "B 4 DH DW"]:
+        self.scheduler.config.num_train_timesteps = t.item()
+        self.scheduler.set_timesteps(self.cfg.diffusion_steps)
+        if mask is not None:
+            mask = F.interpolate(mask, (latents.shape[-2], latents.shape[-1]), mode='bilinear')
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)
+            latents = self.scheduler.add_noise(latents, noise, t)  # type: ignore
+
+            # sections of code used from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
+            threestudio.debug("Start editing...")
+            for i, t in enumerate(self.scheduler.timesteps):
+                # predict the noise residual with unet, NO grad!
+                with torch.no_grad():
+                    # pred noise
+                    latent_model_input = torch.cat([latents] * 2)
+                    (
+                        down_block_res_samples,
+                        mid_block_res_sample,
+                    ) = self.forward_controlnet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=text_embeddings,
+                        image_cond=image_cond,
+                        condition_scale=self.cfg.condition_scale,
+                    )
+
+                    noise_pred = self.forward_control_unet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=text_embeddings,
+                        cross_attention_kwargs=None,
+                        down_block_additional_residuals=down_block_res_samples,
+                        mid_block_additional_residual=mid_block_res_sample,
+                    )
+                # perform classifier-free guidance
+                noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+                if mask is not None:
+                    noise_pred = mask * noise_pred + (1 - mask) * noise
+                # get previous sample, continue loop
+                latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+            threestudio.debug("Editing finished.")
+        return latents
+
+    def prepare_image_cond(self, cond_rgb: Float[Tensor, "B H W C"]):
+        if self.cfg.control_type == "normal":
+            cond_rgb = (
+                (cond_rgb[0].detach().cpu().numpy() * 255).astype(np.uint8).copy()
+            )
+            detected_map = self.preprocessor(cond_rgb)
+            control = (
+                torch.from_numpy(np.array(detected_map)).float().to(self.device) / 255.0
+            )
+            control = control.unsqueeze(0)
+            control = control.permute(0, 3, 1, 2)
+        elif self.cfg.control_type == "canny":
+            cond_rgb = (
+                (cond_rgb[0].detach().cpu().numpy() * 255).astype(np.uint8).copy()
+            )
+            blurred_img = cv2.blur(cond_rgb, ksize=(5, 5))
+            detected_map = self.preprocessor(
+                blurred_img, self.cfg.canny_lower_bound, self.cfg.canny_upper_bound
+            )
+            control = (
+                torch.from_numpy(np.array(detected_map)).float().to(self.device) / 255.0
+            )
+            # control = control.unsqueeze(-1).repeat(1, 1, 3)
+            control = control.unsqueeze(0)
+            control = control.permute(0, 3, 1, 2)
+        elif self.cfg.control_type == "input_normal":
+            cond_rgb[..., 0] = (
+                1 - cond_rgb[..., 0]
+            )  # Flip the sign on the x-axis to match bae system
+            control = cond_rgb.permute(0, 3, 1, 2)
+        else:
+            raise ValueError(f"Unknown control type: {self.cfg.control_type}")
+
+        return control
+
+    def compute_grad_sds(
+        self,
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        latents: Float[Tensor, "B 4 DH DW"],
+        image_cond: Float[Tensor, "B 3 H W"],
+        t: Int[Tensor, "B"],
+    ):
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)  # TODO: use torch generator
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2)
+            down_block_res_samples, mid_block_res_sample = self.forward_controlnet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                image_cond=image_cond,
+                condition_scale=self.cfg.condition_scale,
+            )
+
+            noise_pred = self.forward_control_unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                cross_attention_kwargs=None,
+                down_block_additional_residuals=down_block_res_samples,
+                mid_block_additional_residual=mid_block_res_sample,
+            )
+
+        # perform classifier-free guidance
+        noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_text - noise_pred_uncond
+        )
+
+        w = (1 - self.alphas[t]).view(-1, 1, 1, 1)
+        grad = w * (noise_pred - noise)
+        return grad
+
+    def compute_grad_du(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+        rgb_BCHW_HW8: Float[Tensor, "B 3 RH RW"],
+        cond_feature: Float[Tensor, "B 3 RH RW"],
+        cond_rgb: Float[Tensor, "B H W 3"],
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        mask = None,
+        **kwargs,
+    ):
+        batch_size, _, RH, RW = cond_feature.shape
+        assert batch_size == 1
+
+        origin_gt_rgb = F.interpolate(
+            cond_rgb.permute(0, 3, 1, 2), (RH, RW), mode="bilinear"
+        ).permute(0, 2, 3, 1)
+        need_diffusion = (
+            self.global_step % self.cfg.per_du_step == 0
+            and self.global_step > self.cfg.start_du_step
+        )
+        if self.cfg.cache_du:
+            if torch.is_tensor(kwargs["index"]):
+                batch_index = kwargs["index"].item()
+            else:
+                batch_index = kwargs["index"]
+            if (
+                not (batch_index in self.edit_frames)
+            ) and self.global_step > self.cfg.start_du_step:
+                need_diffusion = True
+        need_loss = self.cfg.cache_du or need_diffusion
+        guidance_out = {}
+
+        if need_diffusion:
+            t = torch.randint(
+                self.min_step,
+                self.max_step,
+                [1],
+                dtype=torch.long,
+                device=self.device,
+            )
+            print("t:", t)
+            edit_latents = self.edit_latents(text_embeddings, latents, cond_feature, t, mask)
+            edit_images = self.decode_latents(edit_latents)
+            edit_images = F.interpolate(
+                edit_images, (RH, RW), mode="bilinear"
+            ).permute(0, 2, 3, 1)
+            self.edit_images = edit_images
+            if self.cfg.cache_du:
+                self.edit_frames[batch_index] = edit_images.detach().cpu()
+
+        if need_loss:
+            if self.cfg.cache_du:
+                if batch_index in self.edit_frames:
+                    gt_rgb = self.edit_frames[batch_index].to(cond_feature.device)
+                else:
+                    gt_rgb = origin_gt_rgb
+            else:
+                gt_rgb = edit_images
+
+            import cv2
+            import numpy as np
+
+            temp = (edit_images.detach().cpu()[0].numpy() * 255).astype(np.uint8)
+            cv2.imwrite(".threestudio_cache/test.jpg", temp[:, :, ::-1])
+            
+            guidance_out.update(
+                {
+                    "loss_l1": torch.nn.functional.l1_loss(
+                        rgb_BCHW_HW8, gt_rgb.permute(0, 3, 1, 2), reduction="sum"
+                    ),
+                    "loss_p": self.perceptual_loss(
+                        rgb_BCHW_HW8.contiguous(),
+                        gt_rgb.permute(0, 3, 1, 2).contiguous(),
+                    ).sum(),
+                }
+            )
+
+        return guidance_out
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        cond_rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        mask = None,
+        **kwargs,
+    ):
+        batch_size, H, W, _ = rgb.shape
+        assert batch_size == 1
+        assert rgb.shape[:-1] == cond_rgb.shape[:-1]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        if mask is not None: mask = mask.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 DH DW"]
+        if self.cfg.fixed_size > 0:
+            RH, RW = self.cfg.fixed_size, self.cfg.fixed_size
+        else:
+            RH, RW = H // 8 * 8, W // 8 * 8
+        rgb_BCHW_HW8 = F.interpolate(
+            rgb_BCHW, (RH, RW), mode="bilinear", align_corners=False
+        )
+        latents = self.encode_images(rgb_BCHW_HW8)
+
+        image_cond = self.prepare_image_cond(cond_rgb)
+        image_cond = F.interpolate(
+            image_cond, (RH, RW), mode="bilinear", align_corners=False
+        )
+
+        temp = torch.zeros(1).to(rgb.device)
+        azimuth = kwargs.get("azimuth", temp)
+        camera_distance = kwargs.get("camera_distance", temp)
+        view_dependent_prompt = kwargs.get("view_dependent_prompt", False)
+        text_embeddings = prompt_utils.get_text_embeddings(temp, azimuth, camera_distance, view_dependent_prompt) # FIXME: change to view-conditioned prompt
+
+        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [batch_size],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+
+        guidance_out = {}
+        if self.cfg.use_sds:
+            grad = self.compute_grad_sds(text_embeddings, latents, image_cond, t)
+            grad = torch.nan_to_num(grad)
+            if self.grad_clip_val is not None:
+                grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+            target = (latents - grad).detach()
+            loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+            guidance_out.update(
+                {
+                    "loss_sds": loss_sds,
+                    "grad_norm": grad.norm(),
+                    "min_step": self.min_step,
+                    "max_step": self.max_step,
+                }
+            )
+
+        if self.cfg.use_du:
+            grad = self.compute_grad_du(
+                latents, rgb_BCHW_HW8, image_cond, cond_rgb, text_embeddings, mask, **kwargs
+            )
+            guidance_out.update(grad)
+
+        return guidance_out
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )
+
+        self.global_step = global_step

threestudio/models/guidance/controlnet_reg_guidance.py

@@ -0,0 +1,454 @@
+import os
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from controlnet_aux import CannyDetector, NormalBaeDetector
+from diffusers import ControlNetModel, DDIMScheduler, StableDiffusionControlNetPipeline, DPMSolverMultistepScheduler
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.typing import *
+
+
+@threestudio.register("stable-diffusion-controlnet-reg-guidance")
+class ControlNetGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+        pretrained_model_name_or_path: str = "SG161222/Realistic_Vision_V2.0"
+        ddim_scheduler_name_or_path: str = "runwayml/stable-diffusion-v1-5"
+        control_type: str = "normal"  # normal/canny
+
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = False
+        guidance_scale: float = 7.5
+        condition_scale: float = 1.5
+        grad_clip: Optional[Any] = None
+        half_precision_weights: bool = True
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+        diffusion_steps: int = 20
+
+        use_sds: bool = False
+
+        # Canny threshold
+        canny_lower_bound: int = 50
+        canny_upper_bound: int = 100
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading ControlNet ...")
+
+        self.weights_dtype = torch.float16 if self.cfg.half_precision_weights else torch.float32
+
+        self.preprocessor, controlnet_name_or_path = self.get_preprocessor_and_controlnet()
+        
+        pipe_kwargs = self.configure_pipeline()
+
+        self.load_models(pipe_kwargs, controlnet_name_or_path)
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.scheduler.set_timesteps(self.cfg.diffusion_steps)
+        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
+        self.scheduler = self.pipe.scheduler
+
+        self.check_memory_efficiency_conditions()
+
+        self.set_min_max_steps()
+        self.alphas = self.scheduler.alphas_cumprod.to(self.device)
+        self.grad_clip_val = None
+
+        threestudio.info(f"Loaded ControlNet!")
+
+    def get_preprocessor_and_controlnet(self):
+        if self.cfg.control_type in ("normal", "input_normal"):
+            if self.cfg.pretrained_model_name_or_path == "SG161222/Realistic_Vision_V2.0":
+                controlnet_name_or_path = "lllyasviel/control_v11p_sd15_normalbae"
+            else:
+                controlnet_name_or_path = "thibaud/controlnet-sd21-normalbae-diffusers"
+            preprocessor = NormalBaeDetector.from_pretrained("lllyasviel/Annotators", cache_dir=self.cfg.cache_dir)
+            preprocessor.model.to(self.device)
+        elif self.cfg.control_type == "canny" or self.cfg.control_type == "canny2":
+            controlnet_name_or_path = self.get_canny_controlnet()
+            preprocessor = CannyDetector()
+        else:
+            raise ValueError(f"Unknown control type: {self.cfg.control_type}")
+        return preprocessor, controlnet_name_or_path
+
+    def get_canny_controlnet(self):
+        if self.cfg.control_type == "canny":
+            return "lllyasviel/control_v11p_sd15_canny"
+        elif self.cfg.control_type == "canny2":
+            return "thepowefuldeez/sd21-controlnet-canny"
+
+    def configure_pipeline(self):
+        return {
+            "safety_checker": None,
+            "feature_extractor": None,
+            "requires_safety_checker": False,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+            "local_files_only": self.cfg.local_files_only
+        }
+
+    def load_models(self, pipe_kwargs, controlnet_name_or_path):
+        controlnet = ControlNetModel.from_pretrained(
+            controlnet_name_or_path,
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+            local_files_only=self.cfg.local_files_only
+        )
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path, controlnet=controlnet, **pipe_kwargs
+        ).to(self.device)
+
+        self.scheduler = DDIMScheduler.from_pretrained(
+            self.cfg.ddim_scheduler_name_or_path,
+            subfolder="scheduler",
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+            local_files_only=self.cfg.local_files_only
+        )
+
+        self.vae = self.pipe.vae.eval()
+        self.unet = self.pipe.unet.eval()
+        self.controlnet = self.pipe.controlnet.eval()
+
+    def check_memory_efficiency_conditions(self):
+        if self.cfg.enable_memory_efficient_attention:
+            self.memory_efficiency_status()
+        if self.cfg.enable_sequential_cpu_offload:
+            self.pipe.enable_sequential_cpu_offload()
+        if self.cfg.enable_attention_slicing:
+            self.pipe.enable_attention_slicing(1)
+        if self.cfg.enable_channels_last_format:
+            self.pipe.unet.to(memory_format=torch.channels_last)
+
+    def memory_efficiency_status(self):
+        if parse_version(torch.__version__) >= parse_version("2"):
+            threestudio.info("PyTorch2.0 uses memory efficient attention by default.")
+        elif not is_xformers_available():
+            threestudio.warn("xformers is not available, memory efficient attention is not enabled.")
+        else:
+            self.pipe.enable_xformers_memory_efficient_attention()
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_controlnet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        image_cond: Float[Tensor, "..."],
+        condition_scale: float,
+        encoder_hidden_states: Float[Tensor, "..."],
+    ) -> Float[Tensor, "..."]:
+        return self.controlnet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+            controlnet_cond=image_cond.to(self.weights_dtype),
+            conditioning_scale=condition_scale,
+            return_dict=False,
+        )
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_control_unet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+        cross_attention_kwargs,
+        down_block_additional_residuals,
+        mid_block_additional_residual,
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        return self.unet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+        ).sample.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 512 512"]
+    ) -> Float[Tensor, "B 4 64 64"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_cond_images(
+        self, imgs: Float[Tensor, "B 3 512 512"]
+    ) -> Float[Tensor, "B 4 64 64"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.mode()
+        uncond_image_latents = torch.zeros_like(latents)
+        latents = torch.cat([latents, latents, uncond_image_latents], dim=0)
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+        latent_height: int = 64,
+        latent_width: int = 64,
+    ) -> Float[Tensor, "B 3 512 512"]:
+        input_dtype = latents.dtype
+        latents = F.interpolate(
+            latents, (latent_height, latent_width), mode="bilinear", align_corners=False
+        )
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents.to(self.weights_dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    def edit_latents(
+        self,
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        latents: Float[Tensor, "B 4 64 64"],
+        image_cond: Float[Tensor, "B 3 512 512"],
+        t: Int[Tensor, "B"],
+        mask=None
+    ) -> Float[Tensor, "B 4 64 64"]:
+        batch_size = t.shape[0]
+        self.scheduler.set_timesteps(num_inference_steps=self.cfg.diffusion_steps)
+        init_timestep = max(1, min(int(self.cfg.diffusion_steps * t[0].item() / self.num_train_timesteps), self.cfg.diffusion_steps))
+        t_start = max(self.cfg.diffusion_steps - init_timestep, 0)
+        latent_timestep = self.scheduler.timesteps[t_start : t_start + 1].repeat(batch_size)
+        B, _, DH, DW = latents.shape
+        origin_latents = latents.clone()
+        if mask is not None:
+            mask = F.interpolate(mask, (DH, DW), mode="bilinear", antialias=True)
+        
+        with torch.no_grad():
+            # sections of code used from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
+            noise = torch.randn_like(latents)
+            latents = self.scheduler.add_noise(latents, noise, latent_timestep)  # type: ignore
+            threestudio.debug("Start editing...")
+            for i, step in enumerate(range(t_start, self.cfg.diffusion_steps)):
+                timestep = self.scheduler.timesteps[step]
+                # predict the noise residual with unet, NO grad!
+                with torch.no_grad():
+                    # pred noise
+                    latent_model_input = torch.cat([latents] * 2)
+                    (
+                        down_block_res_samples,
+                        mid_block_res_sample,
+                    ) = self.forward_controlnet(
+                        latent_model_input,
+                        timestep,
+                        encoder_hidden_states=text_embeddings,
+                        image_cond=image_cond,
+                        condition_scale=self.cfg.condition_scale,
+                    )
+
+                    noise_pred = self.forward_control_unet(
+                        latent_model_input,
+                        timestep,
+                        encoder_hidden_states=text_embeddings,
+                        cross_attention_kwargs=None,
+                        down_block_additional_residuals=down_block_res_samples,
+                        mid_block_additional_residual=mid_block_res_sample,
+                    )
+                # perform classifier-free guidance
+                noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+
+                if mask is not None:
+                    noise_pred = noise_pred * mask + (1-mask) * noise
+                
+                latents = self.scheduler.step(noise_pred, timestep, latents).prev_sample
+            threestudio.debug("Editing finished.")
+        return latents
+
+    def prepare_image_cond(self, cond_rgb: Float[Tensor, "B H W C"]):
+        if self.cfg.control_type == "normal":
+            cond_rgb = (
+                (cond_rgb[0].detach().cpu().numpy() * 255).astype(np.uint8).copy()
+            )
+            detected_map = self.preprocessor(cond_rgb)
+            control = (
+                torch.from_numpy(np.array(detected_map)).float().to(self.device) / 255.0
+            )
+            control = control.unsqueeze(0)
+            control = control.permute(0, 3, 1, 2)
+        elif self.cfg.control_type == "canny" or self.cfg.control_type == "canny2":
+            cond_rgb = (
+                (cond_rgb[0].detach().cpu().numpy() * 255).astype(np.uint8).copy()
+            )
+            blurred_img = cv2.blur(cond_rgb, ksize=(5, 5))
+            detected_map = self.preprocessor(
+                blurred_img, self.cfg.canny_lower_bound, self.cfg.canny_upper_bound
+            )
+            control = (
+                torch.from_numpy(np.array(detected_map)).float().to(self.device) / 255.0
+            )
+            control = control.unsqueeze(-1).repeat(1, 1, 3)
+            control = control.unsqueeze(0)
+            control = control.permute(0, 3, 1, 2)
+        elif self.cfg.control_type == "input_normal":
+            cond_rgb[..., 0] = (
+                1 - cond_rgb[..., 0]
+            )  # Flip the sign on the x-axis to match bae system
+            control = cond_rgb.permute(0, 3, 1, 2)
+        else:
+            raise ValueError(f"Unknown control type: {self.cfg.control_type}")
+
+        return F.interpolate(control, (512, 512), mode="bilinear", align_corners=False)
+
+    def compute_grad_sds(
+        self,
+        text_embeddings: Float[Tensor, "BB 77 768"],
+        latents: Float[Tensor, "B 4 64 64"],
+        image_cond: Float[Tensor, "B 3 512 512"],
+        t: Int[Tensor, "B"],
+    ):
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)  # TODO: use torch generator
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2)
+            down_block_res_samples, mid_block_res_sample = self.forward_controlnet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                image_cond=image_cond,
+                condition_scale=self.cfg.condition_scale,
+            )
+
+            noise_pred = self.forward_control_unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                cross_attention_kwargs=None,
+                down_block_additional_residuals=down_block_res_samples,
+                mid_block_additional_residual=mid_block_res_sample,
+            )
+
+        # perform classifier-free guidance
+        noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_text - noise_pred_uncond
+        )
+
+        w = (1 - self.alphas[t]).view(-1, 1, 1, 1)
+        grad = w * (noise_pred - noise)
+        return grad
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        cond_rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        mask: Float[Tensor, "B H W C"],
+        **kwargs,
+    ):
+        
+        batch_size, H, W, _ = rgb.shape
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 64 64"]
+        rgb_BCHW_512 = F.interpolate(
+            rgb_BCHW, (512, 512), mode="bilinear", align_corners=False
+        )
+        latents = self.encode_images(rgb_BCHW_512)
+
+        image_cond = self.prepare_image_cond(cond_rgb)
+
+        temp = torch.zeros(1).to(rgb.device)
+        text_embeddings = prompt_utils.get_text_embeddings(temp, temp, temp, False)
+
+        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [batch_size],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        if self.cfg.use_sds:
+            grad = self.compute_grad_sds(text_embeddings, latents, image_cond, t)
+            grad = torch.nan_to_num(grad)
+            if self.grad_clip_val is not None:
+                grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+            target = (latents - grad).detach()
+            loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+            return {
+                "loss_sds": loss_sds,
+                "grad_norm": grad.norm(),
+                "min_step": self.min_step,
+                "max_step": self.max_step,
+            }
+        else:
+            
+            if mask is not None: mask = mask.permute(0, 3, 1, 2)
+            edit_latents = self.edit_latents(text_embeddings, latents, image_cond, t, mask)
+            edit_images = self.decode_latents(edit_latents)
+            edit_images = F.interpolate(edit_images, (H, W), mode="bilinear")
+
+            return {"edit_images": edit_images.permute(0, 2, 3, 1),
+                    "edit_latents": edit_latents}
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )
+
+
+if __name__ == "__main__":
+    from threestudio.utils.config import ExperimentConfig, load_config
+    from threestudio.utils.typing import Optional
+
+    cfg = load_config("configs/experimental/controlnet-normal.yaml")
+    guidance = threestudio.find(cfg.system.guidance_type)(cfg.system.guidance)
+    prompt_processor = threestudio.find(cfg.system.prompt_processor_type)(
+        cfg.system.prompt_processor
+    )
+
+    rgb_image = cv2.imread("assets/face.jpg")[:, :, ::-1].copy() / 255
+    rgb_image = cv2.resize(rgb_image, (512, 512))
+    rgb_image = torch.FloatTensor(rgb_image).unsqueeze(0).to(guidance.device)
+    prompt_utils = prompt_processor()
+    guidance_out = guidance(rgb_image, rgb_image, prompt_utils)
+    edit_image = (
+        (guidance_out["edit_images"][0].detach().cpu().clip(0, 1).numpy() * 255)
+        .astype(np.uint8)[:, :, ::-1]
+        .copy()
+    )
+    os.makedirs(".threestudio_cache", exist_ok=True)
+    cv2.imwrite(".threestudio_cache/edit_image.jpg", edit_image)

threestudio/models/guidance/deep_floyd_guidance.py

@@ -0,0 +1,582 @@
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import IFPipeline, DDPMScheduler
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.ops import perpendicular_component
+from threestudio.utils.typing import *
+
+
+@threestudio.register("deep-floyd-guidance")
+class DeepFloydGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+        pretrained_model_name_or_path: str = "DeepFloyd/IF-I-XL-v1.0"
+        # FIXME: xformers error
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = True
+        guidance_scale: float = 20.0
+        grad_clip: Optional[
+            Any
+        ] = None  # field(default_factory=lambda: [0, 2.0, 8.0, 1000])
+        time_prior: Optional[Any] = None # [w1,w2,s1,s2]
+        half_precision_weights: bool = True
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+        weighting_strategy: str = "sds"
+
+        view_dependent_prompting: bool = True
+
+        """Maximum number of batch items to evaluate guidance for (for debugging) and to save on disk. -1 means save all items."""
+        max_items_eval: int = 4
+
+        lora_weights_path: Optional[str] = None
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading Deep Floyd ...")
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        # Create model
+        self.pipe = IFPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            text_encoder=None,
+            safety_checker=None,
+            watermarker=None,
+            feature_extractor=None,
+            requires_safety_checker=False,
+            variant="fp16" if self.cfg.half_precision_weights else None,
+            torch_dtype=self.weights_dtype,
+            cache_dir=self.cfg.cache_dir,
+            local_files_only=self.cfg.local_files_only
+        ).to(self.device)
+
+        # Load lora weights
+        if self.cfg.lora_weights_path is not None:
+            self.pipe.load_lora_weights(self.cfg.lora_weights_path)
+            self.pipe.scheduler = self.pipe.scheduler.__class__.from_config(self.pipe.scheduler.config, variance_type="fixed_small")
+
+        if self.cfg.enable_memory_efficient_attention:
+            if parse_version(torch.__version__) >= parse_version("2"):
+                threestudio.info(
+                    "PyTorch2.0 uses memory efficient attention by default."
+                )
+            elif not is_xformers_available():
+                threestudio.warn(
+                    "xformers is not available, memory efficient attention is not enabled."
+                )
+            else:
+                threestudio.warn(
+                    f"Use DeepFloyd with xformers may raise error, see https://github.com/deep-floyd/IF/issues/52 to track this problem."
+                )
+                self.pipe.enable_xformers_memory_efficient_attention()
+
+        if self.cfg.enable_sequential_cpu_offload:
+            self.pipe.enable_sequential_cpu_offload()
+
+        if self.cfg.enable_attention_slicing:
+            self.pipe.enable_attention_slicing(1)
+
+        if self.cfg.enable_channels_last_format:
+            self.pipe.unet.to(memory_format=torch.channels_last)
+
+        self.unet = self.pipe.unet.eval()
+
+        for p in self.unet.parameters():
+            p.requires_grad_(False)
+
+        self.scheduler = self.pipe.scheduler
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+        if self.cfg.time_prior is not None:
+            m1, m2, s1, s2 = self.cfg.time_prior
+            weights = torch.cat(
+                (
+                    torch.exp(
+                        -((torch.arange(self.num_train_timesteps, m1, -1) - m1) ** 2)
+                        / (2 * s1**2)
+                    ),
+                    torch.ones(m1 - m2 + 1),
+                    torch.exp(
+                        -((torch.arange(m2 - 1, 0, -1) - m2) ** 2) / (2 * s2**2)
+                    ),
+                )
+            )
+            weights = weights / torch.sum(weights)
+            self.time_prior_acc_weights = torch.cumsum(weights, dim=0)
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.grad_clip_val: Optional[float] = None
+
+        threestudio.info(f"Loaded Deep Floyd!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_unet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        return self.unet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+        ).sample.to(input_dtype)
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        current_step_ratio=None,
+        mask: Float[Tensor, "B H W 1"] = None,
+        rgb_as_latents=False,
+        guidance_eval=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+
+        if mask is not None: 
+            mask = mask.permute(0, 3, 1, 2)
+            mask = F.interpolate(
+                mask, (64, 64), mode="bilinear", align_corners=False
+            )
+
+        assert rgb_as_latents == False, f"No latent space in {self.__class__.__name__}"
+        rgb_BCHW = rgb_BCHW * 2.0 - 1.0  # scale to [-1, 1] to match the diffusion range
+        latents = F.interpolate(
+            rgb_BCHW, (64, 64), mode="bilinear", align_corners=False
+        )
+
+        if self.cfg.time_prior is not None:
+            time_index = torch.where(
+                (self.time_prior_acc_weights - current_step_ratio) > 0
+            )[0][0]
+            if time_index == 0 or torch.abs(
+                self.time_prior_acc_weights[time_index] - current_step_ratio
+            ) < torch.abs(
+                self.time_prior_acc_weights[time_index - 1] - current_step_ratio
+            ):
+                t = self.num_train_timesteps - time_index
+            else:
+                t = self.num_train_timesteps - time_index + 1
+            t = torch.clip(t, self.min_step, self.max_step + 1)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+
+        else:
+            # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+            t = torch.randint(
+                self.min_step,
+                self.max_step + 1,
+                [batch_size],
+                dtype=torch.long,
+                device=self.device,
+            )
+
+        if prompt_utils.use_perp_neg:
+            (
+                text_embeddings,
+                neg_guidance_weights,
+            ) = prompt_utils.get_text_embeddings_perp_neg(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                noise = torch.randn_like(latents)
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)
+                if mask is not None:
+                    latents_noisy = (1 - mask) * latents + mask * latents_noisy
+                latent_model_input = torch.cat([latents_noisy] * 4, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 4),
+                    encoder_hidden_states=text_embeddings,
+                )  # (4B, 6, 64, 64)
+
+            noise_pred_text, _ = noise_pred[:batch_size].split(3, dim=1)
+            noise_pred_uncond, _ = noise_pred[batch_size : batch_size * 2].split(
+                3, dim=1
+            )
+            noise_pred_neg, _ = noise_pred[batch_size * 2 :].split(3, dim=1)
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            neg_guidance_weights = None
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                # add noise
+                noise = torch.randn_like(latents)  # TODO: use torch generator
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)
+                if mask is not None:
+                    latents_noisy = (1 - mask) * latents + mask * latents_noisy
+                # pred noise
+                latent_model_input = torch.cat([latents_noisy] * 2, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 2),
+                    encoder_hidden_states=text_embeddings,
+                )  # (2B, 6, 64, 64)
+
+            # perform guidance (high scale from paper!)
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred_text, predicted_variance = noise_pred_text.split(3, dim=1)
+            noise_pred_uncond, _ = noise_pred_uncond.split(3, dim=1)
+            noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        """
+        # thresholding, experimental
+        if self.cfg.thresholding:
+            assert batch_size == 1
+            noise_pred = torch.cat([noise_pred, predicted_variance], dim=1)
+            noise_pred = custom_ddpm_step(self.scheduler,
+                noise_pred, int(t.item()), latents_noisy, **self.pipe.prepare_extra_step_kwargs(None, 0.0)
+            )
+        """
+
+        if self.cfg.weighting_strategy == "sds":
+            # w(t), sigma_t^2
+            w = (1 - self.alphas[t]).view(-1, 1, 1, 1)
+        elif self.cfg.weighting_strategy == "uniform":
+            w = 1
+        elif self.cfg.weighting_strategy == "fantasia3d":
+            w = (self.alphas[t] ** 0.5 * (1 - self.alphas[t])).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(
+                f"Unknown weighting strategy: {self.cfg.weighting_strategy}"
+            )
+
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+        # clip grad for stable training?
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # loss = SpecifyGradient.apply(latents, grad)
+        # SpecifyGradient is not straghtforward, use a reparameterization trick instead
+        target = (latents - grad).detach()
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        loss_sd = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sd,
+            "grad_norm": grad.norm(),
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+        }
+
+        # # FIXME: Visualize inpainting results
+        # self.scheduler.set_timesteps(20)
+        # latents = latents_noisy
+        # for t in tqdm(self.scheduler.timesteps):
+        #         # pred noise
+        #         noise_pred = self.get_noise_pred(
+        #             latents, t, text_embeddings, prompt_utils.use_perp_neg, None
+        #         )
+        #         # get prev latent
+        #         prev_latents = latents
+        #         latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"]
+        #         if mask is not None:
+        #             latents = (1 - mask) * prev_latents + mask * latents
+
+        # denoised_img = (latents / 2 + 0.5).permute(0, 2, 3, 1)
+        # guidance_out.update(
+        #     {"denoised_img": denoised_img}
+        # )
+
+        if guidance_eval:
+            guidance_eval_utils = {
+                "use_perp_neg": prompt_utils.use_perp_neg,
+                "neg_guidance_weights": neg_guidance_weights,
+                "text_embeddings": text_embeddings,
+                "t_orig": t,
+                "latents_noisy": latents_noisy,
+                "noise_pred": torch.cat([noise_pred, predicted_variance], dim=1),
+            }
+            guidance_eval_out = self.guidance_eval(**guidance_eval_utils)
+            texts = []
+            for n, e, a, c in zip(
+                guidance_eval_out["noise_levels"], elevation, azimuth, camera_distances
+            ):
+                texts.append(
+                    f"n{n:.02f}\ne{e.item():.01f}\na{a.item():.01f}\nc{c.item():.02f}"
+                )
+            guidance_eval_out.update({"texts": texts})
+            guidance_out.update({"eval": guidance_eval_out})
+
+        return guidance_out
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_noise_pred(
+        self,
+        latents_noisy,
+        t,
+        text_embeddings,
+        use_perp_neg=False,
+        neg_guidance_weights=None,
+    ):
+        batch_size = latents_noisy.shape[0]
+        if use_perp_neg:
+            latent_model_input = torch.cat([latents_noisy] * 4, dim=0)
+            noise_pred = self.forward_unet(
+                latent_model_input,
+                torch.cat([t.reshape(1)] * 4).to(self.device),
+                encoder_hidden_states=text_embeddings,
+            )  # (4B, 6, 64, 64)
+
+            noise_pred_text, _ = noise_pred[:batch_size].split(3, dim=1)
+            noise_pred_uncond, _ = noise_pred[batch_size : batch_size * 2].split(
+                3, dim=1
+            )
+            noise_pred_neg, _ = noise_pred[batch_size * 2 :].split(3, dim=1)
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            latent_model_input = torch.cat([latents_noisy] * 2, dim=0)
+            noise_pred = self.forward_unet(
+                latent_model_input,
+                torch.cat([t.reshape(1)] * 2).to(self.device),
+                encoder_hidden_states=text_embeddings,
+            )  # (2B, 6, 64, 64)
+
+            # perform guidance (high scale from paper!)
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred_text, predicted_variance = noise_pred_text.split(3, dim=1)
+            noise_pred_uncond, _ = noise_pred_uncond.split(3, dim=1)
+            noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        return torch.cat([noise_pred, predicted_variance], dim=1)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def guidance_eval(
+        self,
+        t_orig,
+        text_embeddings,
+        latents_noisy,
+        noise_pred,
+        use_perp_neg=False,
+        neg_guidance_weights=None,
+    ):
+        # use only 50 timesteps, and find nearest of those to t
+        self.scheduler.set_timesteps(50)
+        self.scheduler.timesteps_gpu = self.scheduler.timesteps.to(self.device)
+        bs = (
+            min(self.cfg.max_items_eval, latents_noisy.shape[0])
+            if self.cfg.max_items_eval > 0
+            else latents_noisy.shape[0]
+        )  # batch size
+        large_enough_idxs = self.scheduler.timesteps_gpu.expand([bs, -1]) > t_orig[
+            :bs
+        ].unsqueeze(
+            -1
+        )  # sized [bs,50] > [bs,1]
+        idxs = torch.min(large_enough_idxs, dim=1)[1]
+        t = self.scheduler.timesteps_gpu[idxs]
+
+        fracs = list((t / self.scheduler.config.num_train_timesteps).cpu().numpy())
+        imgs_noisy = (latents_noisy[:bs] / 2 + 0.5).permute(0, 2, 3, 1)
+
+        # get prev latent
+        latents_1step = []
+        pred_1orig = []
+        for b in range(bs):
+            step_output = self.scheduler.step(
+                noise_pred[b : b + 1], t[b], latents_noisy[b : b + 1]
+            )
+            latents_1step.append(step_output["prev_sample"])
+            pred_1orig.append(step_output["pred_original_sample"])
+        latents_1step = torch.cat(latents_1step)
+        pred_1orig = torch.cat(pred_1orig)
+        imgs_1step = (latents_1step / 2 + 0.5).permute(0, 2, 3, 1)
+        imgs_1orig = (pred_1orig / 2 + 0.5).permute(0, 2, 3, 1)
+
+        latents_final = []
+        for b, i in enumerate(idxs):
+            latents = latents_1step[b : b + 1]
+            text_emb = (
+                text_embeddings[
+                    [b, b + len(idxs), b + 2 * len(idxs), b + 3 * len(idxs)], ...
+                ]
+                if use_perp_neg
+                else text_embeddings[[b, b + len(idxs)], ...]
+            )
+            neg_guid = neg_guidance_weights[b : b + 1] if use_perp_neg else None
+            for t in tqdm(self.scheduler.timesteps[i + 1 :], leave=False):
+                # pred noise
+                noise_pred = self.get_noise_pred(
+                    latents, t, text_emb, use_perp_neg, neg_guid
+                )
+                # get prev latent
+                latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"]
+            latents_final.append(latents)
+
+        latents_final = torch.cat(latents_final)
+        imgs_final = (latents_final / 2 + 0.5).permute(0, 2, 3, 1)
+
+        return {
+            "bs": bs,
+            "noise_levels": fracs,
+            "imgs_noisy": imgs_noisy,
+            "imgs_1step": imgs_1step,
+            "imgs_1orig": imgs_1orig,
+            "imgs_final": imgs_final,
+        }
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )
+
+
+"""
+# used by thresholding, experimental
+def custom_ddpm_step(ddpm, model_output: torch.FloatTensor, timestep: int, sample: torch.FloatTensor, generator=None, return_dict: bool = True):
+    self = ddpm
+    t = timestep
+
+    prev_t = self.previous_timestep(t)
+
+    if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:
+        model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)
+    else:
+        predicted_variance = None
+
+    # 1. compute alphas, betas
+    alpha_prod_t = self.alphas_cumprod[t].item()
+    alpha_prod_t_prev = self.alphas_cumprod[prev_t].item() if prev_t >= 0 else 1.0
+    beta_prod_t = 1 - alpha_prod_t
+    beta_prod_t_prev = 1 - alpha_prod_t_prev
+    current_alpha_t = alpha_prod_t / alpha_prod_t_prev
+    current_beta_t = 1 - current_alpha_t
+
+    # 2. compute predicted original sample from predicted noise also called
+    # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
+    if self.config.prediction_type == "epsilon":
+        pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+    elif self.config.prediction_type == "sample":
+        pred_original_sample = model_output
+    elif self.config.prediction_type == "v_prediction":
+        pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
+    else:
+        raise ValueError(
+            f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
+            " `v_prediction`  for the DDPMScheduler."
+        )
+
+    # 3. Clip or threshold "predicted x_0"
+    if self.config.thresholding:
+        pred_original_sample = self._threshold_sample(pred_original_sample)
+    elif self.config.clip_sample:
+        pred_original_sample = pred_original_sample.clamp(
+            -self.config.clip_sample_range, self.config.clip_sample_range
+        )
+
+    noise_thresholded = (sample - (alpha_prod_t ** 0.5) * pred_original_sample) / (beta_prod_t ** 0.5)
+    return noise_thresholded
+"""
+
+
+if __name__ == '__main__':
+    from threestudio.utils.config import load_config
+    import pytorch_lightning as pl
+    import numpy as np
+    import os
+    import cv2
+    cfg = load_config("configs/debugging/deepfloyd.yaml")
+    guidance = threestudio.find(cfg.system.guidance_type)(cfg.system.guidance)
+    prompt_processor = threestudio.find(cfg.system.prompt_processor_type)(cfg.system.prompt_processor)
+    prompt_utils = prompt_processor()
+    temp = torch.zeros(1).to(guidance.device)
+    # rgb_image = guidance.sample(prompt_utils, temp, temp, temp, seed=cfg.seed)
+    # rgb_image = (rgb_image[0].detach().cpu().clip(0, 1).numpy()*255).astype(np.uint8)[:, :, ::-1].copy()
+    # os.makedirs('.threestudio_cache', exist_ok=True)
+    # cv2.imwrite('.threestudio_cache/diffusion_image.jpg', rgb_image)
+
+    ### inpaint
+    rgb_image = cv2.imread("assets/test.jpg")[:, :, ::-1].copy() / 255
+    mask_image = cv2.imread("assets/mask.png")[:, :, :1].copy() / 255
+    rgb_image = cv2.resize(rgb_image, (512, 512))
+    mask_image = cv2.resize(mask_image, (512, 512)).reshape(512, 512, 1)
+    rgb_image = torch.FloatTensor(rgb_image).unsqueeze(0).to(guidance.device)
+    mask_image = torch.FloatTensor(mask_image).unsqueeze(0).to(guidance.device)
+
+    guidance_out = guidance(rgb_image, prompt_utils, temp, temp, temp, mask=mask_image)
+    edit_image = (
+        (guidance_out["denoised_img"][0].detach().cpu().clip(0, 1).numpy() * 255)
+        .astype(np.uint8)[:, :, ::-1]
+        .copy()
+    )
+    os.makedirs(".threestudio_cache", exist_ok=True)
+    cv2.imwrite(".threestudio_cache/edit_image.jpg", edit_image)

threestudio/models/guidance/stable_diffusion_guidance.py

@@ -0,0 +1,632 @@
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import DDIMScheduler, DDPMScheduler, StableDiffusionPipeline
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, cleanup, parse_version
+from threestudio.utils.ops import perpendicular_component
+from threestudio.utils.typing import *
+
+
+@threestudio.register("stable-diffusion-guidance")
+class StableDiffusionGuidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+        pretrained_model_name_or_path: str = "runwayml/stable-diffusion-v1-5"
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = False
+        guidance_scale: float = 100.0
+        grad_clip: Optional[
+            Any
+        ] = None  # field(default_factory=lambda: [0, 2.0, 8.0, 1000])
+        time_prior: Optional[Any] = None # [w1,w2,s1,s2]
+        half_precision_weights: bool = True
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+        max_step_percent_annealed: float = 0.5
+        anneal_start_step: Optional[int] = None
+
+        use_sjc: bool = False
+        var_red: bool = True
+        weighting_strategy: str = "sds"
+
+        token_merging: bool = False
+        token_merging_params: Optional[dict] = field(default_factory=dict)
+
+        view_dependent_prompting: bool = True
+
+        """Maximum number of batch items to evaluate guidance for (for debugging) and to save on disk. -1 means save all items."""
+        max_items_eval: int = 4
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading Stable Diffusion ...")
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        pipe_kwargs = {
+            "tokenizer": None,
+            "safety_checker": None,
+            "feature_extractor": None,
+            "requires_safety_checker": False,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+            "local_files_only": self.cfg.local_files_only
+        }
+        self.pipe = StableDiffusionPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            **pipe_kwargs,
+        ).to(self.device)
+
+        if self.cfg.enable_memory_efficient_attention:
+            if parse_version(torch.__version__) >= parse_version("2"):
+                threestudio.info(
+                    "PyTorch2.0 uses memory efficient attention by default."
+                )
+            elif not is_xformers_available():
+                threestudio.warn(
+                    "xformers is not available, memory efficient attention is not enabled."
+                )
+            else:
+                self.pipe.enable_xformers_memory_efficient_attention()
+
+        if self.cfg.enable_sequential_cpu_offload:
+            self.pipe.enable_sequential_cpu_offload()
+
+        if self.cfg.enable_attention_slicing:
+            self.pipe.enable_attention_slicing(1)
+
+        if self.cfg.enable_channels_last_format:
+            self.pipe.unet.to(memory_format=torch.channels_last)
+
+        del self.pipe.text_encoder
+        cleanup()
+
+        # Create model
+        self.vae = self.pipe.vae.eval()
+        self.unet = self.pipe.unet.eval()
+
+        for p in self.vae.parameters():
+            p.requires_grad_(False)
+        for p in self.unet.parameters():
+            p.requires_grad_(False)
+
+        if self.cfg.token_merging:
+            import tomesd
+
+            tomesd.apply_patch(self.unet, **self.cfg.token_merging_params)
+
+        if self.cfg.use_sjc:
+            # score jacobian chaining use DDPM
+            self.scheduler = DDPMScheduler.from_pretrained(
+                self.cfg.pretrained_model_name_or_path,
+                subfolder="scheduler",
+                torch_dtype=self.weights_dtype,
+                beta_start=0.00085,
+                beta_end=0.0120,
+                beta_schedule="scaled_linear",
+                cache_dir=self.cfg.cache_dir,
+            )
+        else:
+            self.scheduler = DDIMScheduler.from_pretrained(
+                self.cfg.pretrained_model_name_or_path,
+                subfolder="scheduler",
+                torch_dtype=self.weights_dtype,
+                cache_dir=self.cfg.cache_dir,
+                local_files_only=self.cfg.local_files_only,
+            )
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+        if self.cfg.time_prior is not None:
+            m1, m2, s1, s2 = self.cfg.time_prior
+            weights = torch.cat(
+                (
+                    torch.exp(
+                        -((torch.arange(self.num_train_timesteps, m1, -1) - m1) ** 2)
+                        / (2 * s1**2)
+                    ),
+                    torch.ones(m1 - m2 + 1),
+                    torch.exp(
+                        -((torch.arange(m2 - 1, 0, -1) - m2) ** 2) / (2 * s2**2)
+                    ),
+                )
+            )
+            weights = weights / torch.sum(weights)
+            self.time_prior_acc_weights = torch.cumsum(weights, dim=0)
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+        if self.cfg.use_sjc:
+            # score jacobian chaining need mu
+            self.us: Float[Tensor, "..."] = torch.sqrt((1 - self.alphas) / self.alphas)
+
+        self.grad_clip_val: Optional[float] = None
+
+        threestudio.info(f"Loaded Stable Diffusion!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_unet(
+        self,
+        latents: Float[Tensor, "..."],
+        t: Float[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        return self.unet(
+            latents.to(self.weights_dtype),
+            t.to(self.weights_dtype),
+            encoder_hidden_states=encoder_hidden_states.to(self.weights_dtype),
+        ).sample.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 512 512"]
+    ) -> Float[Tensor, "B 4 64 64"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+        latent_height: int = 64,
+        latent_width: int = 64,
+    ) -> Float[Tensor, "B 3 512 512"]:
+        input_dtype = latents.dtype
+        latents = F.interpolate(
+            latents, (latent_height, latent_width), mode="bilinear", align_corners=False
+        )
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents.to(self.weights_dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    def compute_grad_sds(
+        self,
+        latents: Float[Tensor, "B 4 64 64"],
+        t: Int[Tensor, "B"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ):
+        batch_size = elevation.shape[0]
+
+        if prompt_utils.use_perp_neg:
+            (
+                text_embeddings,
+                neg_guidance_weights,
+            ) = prompt_utils.get_text_embeddings_perp_neg(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                noise = torch.randn_like(latents)
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)
+                latent_model_input = torch.cat([latents_noisy] * 4, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 4),
+                    encoder_hidden_states=text_embeddings,
+                )  # (4B, 3, 64, 64)
+
+            noise_pred_text = noise_pred[:batch_size]
+            noise_pred_uncond = noise_pred[batch_size : batch_size * 2]
+            noise_pred_neg = noise_pred[batch_size * 2 :]
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            neg_guidance_weights = None
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                # add noise
+                noise = torch.randn_like(latents)  # TODO: use torch generator
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)
+                # pred noise
+                latent_model_input = torch.cat([latents_noisy] * 2, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 2),
+                    encoder_hidden_states=text_embeddings,
+                )
+
+            # perform guidance (high scale from paper!)
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        if self.cfg.weighting_strategy == "sds":
+            # w(t), sigma_t^2
+            w = (1 - self.alphas[t]).view(-1, 1, 1, 1)
+        elif self.cfg.weighting_strategy == "uniform":
+            w = 1
+        elif self.cfg.weighting_strategy == "fantasia3d":
+            w = (self.alphas[t] ** 0.5 * (1 - self.alphas[t])).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(
+                f"Unknown weighting strategy: {self.cfg.weighting_strategy}"
+            )
+
+        grad = w * (noise_pred - noise)
+
+        guidance_eval_utils = {
+            "use_perp_neg": prompt_utils.use_perp_neg,
+            "neg_guidance_weights": neg_guidance_weights,
+            "text_embeddings": text_embeddings,
+            "t_orig": t,
+            "latents_noisy": latents_noisy,
+            "noise_pred": noise_pred,
+        }
+
+        return grad, guidance_eval_utils
+
+    def compute_grad_sjc(
+        self,
+        latents: Float[Tensor, "B 4 64 64"],
+        t: Int[Tensor, "B"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ):
+        batch_size = elevation.shape[0]
+
+        sigma = self.us[t]
+        sigma = sigma.view(-1, 1, 1, 1)
+
+        if prompt_utils.use_perp_neg:
+            (
+                text_embeddings,
+                neg_guidance_weights,
+            ) = prompt_utils.get_text_embeddings_perp_neg(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                noise = torch.randn_like(latents)
+                y = latents
+                zs = y + sigma * noise
+                scaled_zs = zs / torch.sqrt(1 + sigma**2)
+                # pred noise
+                latent_model_input = torch.cat([scaled_zs] * 4, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 4),
+                    encoder_hidden_states=text_embeddings,
+                )  # (4B, 3, 64, 64)
+
+            noise_pred_text = noise_pred[:batch_size]
+            noise_pred_uncond = noise_pred[batch_size : batch_size * 2]
+            noise_pred_neg = noise_pred[batch_size * 2 :]
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            neg_guidance_weights = None
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                # add noise
+                noise = torch.randn_like(latents)  # TODO: use torch generator
+                y = latents
+
+                zs = y + sigma * noise
+                scaled_zs = zs / torch.sqrt(1 + sigma**2)
+
+                # pred noise
+                latent_model_input = torch.cat([scaled_zs] * 2, dim=0)
+                noise_pred = self.forward_unet(
+                    latent_model_input,
+                    torch.cat([t] * 2),
+                    encoder_hidden_states=text_embeddings,
+                )
+
+                # perform guidance (high scale from paper!)
+                noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+                noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+
+        Ds = zs - sigma * noise_pred
+
+        if self.cfg.var_red:
+            grad = -(Ds - y) / sigma
+        else:
+            grad = -(Ds - zs) / sigma
+
+        guidance_eval_utils = {
+            "use_perp_neg": prompt_utils.use_perp_neg,
+            "neg_guidance_weights": neg_guidance_weights,
+            "text_embeddings": text_embeddings,
+            "t_orig": t,
+            "latents_noisy": scaled_zs,
+            "noise_pred": noise_pred,
+        }
+
+        return grad, guidance_eval_utils
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        rgb_as_latents=False,
+        guidance_eval=False,
+        current_step_ratio=None,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 64 64"]
+        if rgb_as_latents:
+            latents = F.interpolate(
+                rgb_BCHW, (64, 64), mode="bilinear", align_corners=False
+            )
+        else:
+            rgb_BCHW_512 = F.interpolate(
+                rgb_BCHW, (512, 512), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.encode_images(rgb_BCHW_512)
+
+        if self.cfg.time_prior is not None:
+            time_index = torch.where(
+                (self.time_prior_acc_weights - current_step_ratio) > 0
+            )[0][0]
+            if time_index == 0 or torch.abs(
+                self.time_prior_acc_weights[time_index] - current_step_ratio
+            ) < torch.abs(
+                self.time_prior_acc_weights[time_index - 1] - current_step_ratio
+            ):
+                t = self.num_train_timesteps - time_index
+            else:
+                t = self.num_train_timesteps - time_index + 1
+            t = torch.clip(t, self.min_step, self.max_step + 1)
+            t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
+
+        else:
+            # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+            t = torch.randint(
+                self.min_step,
+                self.max_step + 1,
+                [batch_size],
+                dtype=torch.long,
+                device=self.device,
+            )
+
+        if self.cfg.use_sjc:
+            grad, guidance_eval_utils = self.compute_grad_sjc(
+                latents, t, prompt_utils, elevation, azimuth, camera_distances
+            )
+        else:
+            grad, guidance_eval_utils = self.compute_grad_sds(
+                latents, t, prompt_utils, elevation, azimuth, camera_distances
+            )
+
+        grad = torch.nan_to_num(grad)
+        # clip grad for stable training?
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # loss = SpecifyGradient.apply(latents, grad)
+        # SpecifyGradient is not straghtforward, use a reparameterization trick instead
+        target = (latents - grad).detach()
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sds,
+            "grad_norm": grad.norm(),
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+        }
+
+        if guidance_eval:
+            guidance_eval_out = self.guidance_eval(**guidance_eval_utils)
+            texts = []
+            for n, e, a, c in zip(
+                guidance_eval_out["noise_levels"], elevation, azimuth, camera_distances
+            ):
+                texts.append(
+                    f"n{n:.02f}\ne{e.item():.01f}\na{a.item():.01f}\nc{c.item():.02f}"
+                )
+            guidance_eval_out.update({"texts": texts})
+            guidance_out.update({"eval": guidance_eval_out})
+
+        return guidance_out
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_noise_pred(
+        self,
+        latents_noisy,
+        t,
+        text_embeddings,
+        use_perp_neg=False,
+        neg_guidance_weights=None,
+    ):
+        batch_size = latents_noisy.shape[0]
+
+        if use_perp_neg:
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 4, dim=0)
+            noise_pred = self.forward_unet(
+                latent_model_input,
+                torch.cat([t.reshape(1)] * 4).to(self.device),
+                encoder_hidden_states=text_embeddings,
+            )  # (4B, 3, 64, 64)
+
+            noise_pred_text = noise_pred[:batch_size]
+            noise_pred_uncond = noise_pred[batch_size : batch_size * 2]
+            noise_pred_neg = noise_pred[batch_size * 2 :]
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            # pred noise
+            latent_model_input = torch.cat([latents_noisy] * 2, dim=0)
+            noise_pred = self.forward_unet(
+                latent_model_input,
+                torch.cat([t.reshape(1)] * 2).to(self.device),
+                encoder_hidden_states=text_embeddings,
+            )
+            # perform guidance (high scale from paper!)
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred = noise_pred_text + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        return noise_pred
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def guidance_eval(
+        self,
+        t_orig,
+        text_embeddings,
+        latents_noisy,
+        noise_pred,
+        use_perp_neg=False,
+        neg_guidance_weights=None,
+    ):
+        # use only 50 timesteps, and find nearest of those to t
+        self.scheduler.set_timesteps(50)
+        self.scheduler.timesteps_gpu = self.scheduler.timesteps.to(self.device)
+        bs = (
+            min(self.cfg.max_items_eval, latents_noisy.shape[0])
+            if self.cfg.max_items_eval > 0
+            else latents_noisy.shape[0]
+        )  # batch size
+        large_enough_idxs = self.scheduler.timesteps_gpu.expand([bs, -1]) > t_orig[
+            :bs
+        ].unsqueeze(
+            -1
+        )  # sized [bs,50] > [bs,1]
+        idxs = torch.min(large_enough_idxs, dim=1)[1]
+        t = self.scheduler.timesteps_gpu[idxs]
+
+        fracs = list((t / self.scheduler.config.num_train_timesteps).cpu().numpy())
+        imgs_noisy = self.decode_latents(latents_noisy[:bs]).permute(0, 2, 3, 1)
+
+        # get prev latent
+        latents_1step = []
+        pred_1orig = []
+        for b in range(bs):
+            step_output = self.scheduler.step(
+                noise_pred[b : b + 1], t[b], latents_noisy[b : b + 1], eta=1
+            )
+            latents_1step.append(step_output["prev_sample"])
+            pred_1orig.append(step_output["pred_original_sample"])
+        latents_1step = torch.cat(latents_1step)
+        pred_1orig = torch.cat(pred_1orig)
+        imgs_1step = self.decode_latents(latents_1step).permute(0, 2, 3, 1)
+        imgs_1orig = self.decode_latents(pred_1orig).permute(0, 2, 3, 1)
+
+        latents_final = []
+        for b, i in enumerate(idxs):
+            latents = latents_1step[b : b + 1]
+            text_emb = (
+                text_embeddings[
+                    [b, b + len(idxs), b + 2 * len(idxs), b + 3 * len(idxs)], ...
+                ]
+                if use_perp_neg
+                else text_embeddings[[b, b + len(idxs)], ...]
+            )
+            neg_guid = neg_guidance_weights[b : b + 1] if use_perp_neg else None
+            for t in tqdm(self.scheduler.timesteps[i + 1 :], leave=False):
+                # pred noise
+                noise_pred = self.get_noise_pred(
+                    latents, t, text_emb, use_perp_neg, neg_guid
+                )
+                # get prev latent
+                latents = self.scheduler.step(noise_pred, t, latents, eta=1)[
+                    "prev_sample"
+                ]
+            latents_final.append(latents)
+
+        latents_final = torch.cat(latents_final)
+        imgs_final = self.decode_latents(latents_final).permute(0, 2, 3, 1)
+
+        return {
+            "bs": bs,
+            "noise_levels": fracs,
+            "imgs_noisy": imgs_noisy,
+            "imgs_1step": imgs_1step,
+            "imgs_1orig": imgs_1orig,
+            "imgs_final": imgs_final,
+        }
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )

threestudio/models/guidance/stable_diffusion_unified_guidance.py

@@ -0,0 +1,729 @@
+import random
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import (
+    AutoencoderKL,
+    ControlNetModel,
+    DDPMScheduler,
+    DPMSolverSinglestepScheduler,
+    StableDiffusionPipeline,
+    UNet2DConditionModel,
+)
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import LoRAAttnProcessor
+from diffusers.models.embeddings import TimestepEmbedding
+from diffusers.utils.import_utils import is_xformers_available
+from tqdm import tqdm
+
+import threestudio
+from threestudio.models.networks import ToDTypeWrapper
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseModule
+from threestudio.utils.misc import C, cleanup, enable_gradient, parse_version
+from threestudio.utils.ops import perpendicular_component
+from threestudio.utils.typing import *
+
+
+@threestudio.register("stable-diffusion-unified-guidance")
+class StableDiffusionUnifiedGuidance(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+
+        # guidance type, in ["sds", "vsd"]
+        guidance_type: str = "sds"
+
+        pretrained_model_name_or_path: str = "runwayml/stable-diffusion-v1-5"
+        guidance_scale: float = 100.0
+        weighting_strategy: str = "dreamfusion"
+        view_dependent_prompting: bool = True
+
+        min_step_percent: Any = 0.02
+        max_step_percent: Any = 0.98
+        grad_clip: Optional[Any] = None
+
+        return_rgb_1step_orig: bool = False
+        return_rgb_multistep_orig: bool = False
+        n_rgb_multistep_orig_steps: int = 4
+
+        # TODO
+        # controlnet
+        controlnet_model_name_or_path: Optional[str] = None
+        preprocessor: Optional[str] = None
+        control_scale: float = 1.0
+
+        # TODO
+        # lora
+        lora_model_name_or_path: Optional[str] = None
+
+        # efficiency-related configurations
+        half_precision_weights: bool = True
+        enable_memory_efficient_attention: bool = False
+        enable_sequential_cpu_offload: bool = False
+        enable_attention_slicing: bool = False
+        enable_channels_last_format: bool = False
+        token_merging: bool = False
+        token_merging_params: Optional[dict] = field(default_factory=dict)
+
+        # VSD configurations, only used when guidance_type is "vsd"
+        vsd_phi_model_name_or_path: Optional[str] = None
+        vsd_guidance_scale_phi: float = 1.0
+        vsd_use_lora: bool = True
+        vsd_lora_cfg_training: bool = False
+        vsd_lora_n_timestamp_samples: int = 1
+        vsd_use_camera_condition: bool = True
+        # camera condition type, in ["extrinsics", "mvp", "spherical"]
+        vsd_camera_condition_type: Optional[str] = "extrinsics"
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.min_step: Optional[int] = None
+        self.max_step: Optional[int] = None
+        self.grad_clip_val: Optional[float] = None
+
+        @dataclass
+        class NonTrainableModules:
+            pipe: StableDiffusionPipeline
+            pipe_phi: Optional[StableDiffusionPipeline] = None
+            controlnet: Optional[ControlNetModel] = None
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        threestudio.info(f"Loading Stable Diffusion ...")
+
+        pipe_kwargs = {
+            "tokenizer": None,
+            "safety_checker": None,
+            "feature_extractor": None,
+            "requires_safety_checker": False,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+            "local_files_only": self.cfg.local_files_only,
+        }
+        pipe = StableDiffusionPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            **pipe_kwargs,
+        ).to(self.device)
+        self.prepare_pipe(pipe)
+        self.configure_pipe_token_merging(pipe)
+
+        # phi network for VSD
+        # introduce two trainable modules:
+        # - self.camera_embedding
+        # - self.lora_layers
+        pipe_phi = None
+
+        # if the phi network shares the same unet with the pretrain network
+        # we need to pass additional cross attention kwargs to the unet
+        self.vsd_share_model = (
+            self.cfg.guidance_type == "vsd"
+            and self.cfg.vsd_phi_model_name_or_path is None
+        )
+        if self.cfg.guidance_type == "vsd":
+            if self.cfg.vsd_phi_model_name_or_path is None:
+                pipe_phi = pipe
+            else:
+                pipe_phi = StableDiffusionPipeline.from_pretrained(
+                    self.cfg.vsd_phi_model_name_or_path,
+                    **pipe_kwargs,
+                ).to(self.device)
+                self.prepare_pipe(pipe_phi)
+                self.configure_pipe_token_merging(pipe_phi)
+
+            # set up camera embedding
+            if self.cfg.vsd_use_camera_condition:
+                if self.cfg.vsd_camera_condition_type in ["extrinsics", "mvp"]:
+                    self.camera_embedding_dim = 16
+                elif self.cfg.vsd_camera_condition_type == "spherical":
+                    self.camera_embedding_dim = 4
+                else:
+                    raise ValueError("Invalid camera condition type!")
+
+                # FIXME: hard-coded output dim
+                self.camera_embedding = ToDTypeWrapper(
+                    TimestepEmbedding(self.camera_embedding_dim, 1280),
+                    self.weights_dtype,
+                ).to(self.device)
+                pipe_phi.unet.class_embedding = self.camera_embedding
+
+            if self.cfg.vsd_use_lora:
+                # set up LoRA layers
+                lora_attn_procs = {}
+                for name in pipe_phi.unet.attn_processors.keys():
+                    cross_attention_dim = (
+                        None
+                        if name.endswith("attn1.processor")
+                        else pipe_phi.unet.config.cross_attention_dim
+                    )
+                    if name.startswith("mid_block"):
+                        hidden_size = pipe_phi.unet.config.block_out_channels[-1]
+                    elif name.startswith("up_blocks"):
+                        block_id = int(name[len("up_blocks.")])
+                        hidden_size = list(
+                            reversed(pipe_phi.unet.config.block_out_channels)
+                        )[block_id]
+                    elif name.startswith("down_blocks"):
+                        block_id = int(name[len("down_blocks.")])
+                        hidden_size = pipe_phi.unet.config.block_out_channels[block_id]
+
+                    lora_attn_procs[name] = LoRAAttnProcessor(
+                        hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
+                    )
+
+                pipe_phi.unet.set_attn_processor(lora_attn_procs)
+
+                self.lora_layers = AttnProcsLayers(pipe_phi.unet.attn_processors).to(
+                    self.device
+                )
+                self.lora_layers._load_state_dict_pre_hooks.clear()
+                self.lora_layers._state_dict_hooks.clear()
+
+        threestudio.info(f"Loaded Stable Diffusion!")
+
+        # controlnet
+        controlnet = None
+        if self.cfg.controlnet_model_name_or_path is not None:
+            threestudio.info(f"Loading ControlNet ...")
+
+            controlnet = ControlNetModel.from_pretrained(
+                self.cfg.controlnet_model_name_or_path,
+                torch_dtype=self.weights_dtype,
+            ).to(self.device)
+            controlnet.eval()
+            enable_gradient(controlnet, enabled=False)
+
+            threestudio.info(f"Loaded ControlNet!")
+
+        self.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+
+        # q(z_t|x) = N(alpha_t x, sigma_t^2 I)
+        # in DDPM, alpha_t = sqrt(alphas_cumprod_t), sigma_t^2 = 1 - alphas_cumprod_t
+        self.alphas_cumprod: Float[Tensor, "T"] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+        self.alphas: Float[Tensor, "T"] = self.alphas_cumprod**0.5
+        self.sigmas: Float[Tensor, "T"] = (1 - self.alphas_cumprod) ** 0.5
+        # log SNR
+        self.lambdas: Float[Tensor, "T"] = self.sigmas / self.alphas
+
+        self._non_trainable_modules = NonTrainableModules(
+            pipe=pipe,
+            pipe_phi=pipe_phi,
+            controlnet=controlnet,
+        )
+
+    @property
+    def pipe(self) -> StableDiffusionPipeline:
+        return self._non_trainable_modules.pipe
+
+    @property
+    def pipe_phi(self) -> StableDiffusionPipeline:
+        if self._non_trainable_modules.pipe_phi is None:
+            raise RuntimeError("phi model is not available.")
+        return self._non_trainable_modules.pipe_phi
+
+    @property
+    def controlnet(self) -> ControlNetModel:
+        if self._non_trainable_modules.controlnet is None:
+            raise RuntimeError("ControlNet model is not available.")
+        return self._non_trainable_modules.controlnet
+
+    def prepare_pipe(self, pipe: StableDiffusionPipeline):
+        if self.cfg.enable_memory_efficient_attention:
+            if parse_version(torch.__version__) >= parse_version("2"):
+                threestudio.info(
+                    "PyTorch2.0 uses memory efficient attention by default."
+                )
+            elif not is_xformers_available():
+                threestudio.warn(
+                    "xformers is not available, memory efficient attention is not enabled."
+                )
+            else:
+                pipe.enable_xformers_memory_efficient_attention()
+
+        if self.cfg.enable_sequential_cpu_offload:
+            pipe.enable_sequential_cpu_offload()
+
+        if self.cfg.enable_attention_slicing:
+            pipe.enable_attention_slicing(1)
+
+        if self.cfg.enable_channels_last_format:
+            pipe.unet.to(memory_format=torch.channels_last)
+
+        # FIXME: pipe.__call__ requires text_encoder.dtype
+        # pipe.text_encoder.to("meta")
+        cleanup()
+
+        pipe.vae.eval()
+        pipe.unet.eval()
+
+        enable_gradient(pipe.vae, enabled=False)
+        enable_gradient(pipe.unet, enabled=False)
+
+        # disable progress bar
+        pipe.set_progress_bar_config(disable=True)
+
+    def configure_pipe_token_merging(self, pipe: StableDiffusionPipeline):
+        if self.cfg.token_merging:
+            import tomesd
+
+            tomesd.apply_patch(pipe.unet, **self.cfg.token_merging_params)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_unet(
+        self,
+        unet: UNet2DConditionModel,
+        latents: Float[Tensor, "..."],
+        t: Int[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+        class_labels: Optional[Float[Tensor, "..."]] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        down_block_additional_residuals: Optional[Float[Tensor, "..."]] = None,
+        mid_block_additional_residual: Optional[Float[Tensor, "..."]] = None,
+        velocity_to_epsilon: bool = False,
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        pred = unet(
+            latents.to(unet.dtype),
+            t.to(unet.dtype),
+            encoder_hidden_states=encoder_hidden_states.to(unet.dtype),
+            class_labels=class_labels,
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+        ).sample
+        if velocity_to_epsilon:
+            pred = latents * self.sigmas[t].view(-1, 1, 1, 1) + pred * self.alphas[
+                t
+            ].view(-1, 1, 1, 1)
+        return pred.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def vae_encode(
+        self, vae: AutoencoderKL, imgs: Float[Tensor, "B 3 H W"], mode=False
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        # expect input in [-1, 1]
+        input_dtype = imgs.dtype
+        posterior = vae.encode(imgs.to(vae.dtype)).latent_dist
+        if mode:
+            latents = posterior.mode()
+        else:
+            latents = posterior.sample()
+        latents = latents * vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def vae_decode(
+        self, vae: AutoencoderKL, latents: Float[Tensor, "B 4 Hl Wl"]
+    ) -> Float[Tensor, "B 3 H W"]:
+        # output in [0, 1]
+        input_dtype = latents.dtype
+        latents = 1 / vae.config.scaling_factor * latents
+        image = vae.decode(latents.to(vae.dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    @contextmanager
+    def disable_unet_class_embedding(self, unet: UNet2DConditionModel):
+        class_embedding = unet.class_embedding
+        try:
+            unet.class_embedding = None
+            yield unet
+        finally:
+            unet.class_embedding = class_embedding
+
+    @contextmanager
+    def set_scheduler(
+        self, pipe: StableDiffusionPipeline, scheduler_class: Any, **kwargs
+    ):
+        scheduler_orig = pipe.scheduler
+        pipe.scheduler = scheduler_class.from_config(scheduler_orig.config, **kwargs)
+        yield pipe
+        pipe.scheduler = scheduler_orig
+
+    def get_eps_pretrain(
+        self,
+        latents_noisy: Float[Tensor, "B 4 Hl Wl"],
+        t: Int[Tensor, "B"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        batch_size = latents_noisy.shape[0]
+
+        if prompt_utils.use_perp_neg:
+            (
+                text_embeddings,
+                neg_guidance_weights,
+            ) = prompt_utils.get_text_embeddings_perp_neg(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                with self.disable_unet_class_embedding(self.pipe.unet) as unet:
+                    noise_pred = self.forward_unet(
+                        unet,
+                        torch.cat([latents_noisy] * 4, dim=0),
+                        torch.cat([t] * 4, dim=0),
+                        encoder_hidden_states=text_embeddings,
+                        cross_attention_kwargs={"scale": 0.0}
+                        if self.vsd_share_model
+                        else None,
+                        velocity_to_epsilon=self.pipe.scheduler.config.prediction_type
+                        == "v_prediction",
+                    )  # (4B, 3, Hl, Wl)
+
+            noise_pred_text = noise_pred[:batch_size]
+            noise_pred_uncond = noise_pred[batch_size : batch_size * 2]
+            noise_pred_neg = noise_pred[batch_size * 2 :]
+
+            e_pos = noise_pred_text - noise_pred_uncond
+            accum_grad = 0
+            n_negative_prompts = neg_guidance_weights.shape[-1]
+            for i in range(n_negative_prompts):
+                e_i_neg = noise_pred_neg[i::n_negative_prompts] - noise_pred_uncond
+                accum_grad += neg_guidance_weights[:, i].view(
+                    -1, 1, 1, 1
+                ) * perpendicular_component(e_i_neg, e_pos)
+
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                e_pos + accum_grad
+            )
+        else:
+            text_embeddings = prompt_utils.get_text_embeddings(
+                elevation, azimuth, camera_distances, self.cfg.view_dependent_prompting
+            )
+            with torch.no_grad():
+                with self.disable_unet_class_embedding(self.pipe.unet) as unet:
+                    noise_pred = self.forward_unet(
+                        unet,
+                        torch.cat([latents_noisy] * 2, dim=0),
+                        torch.cat([t] * 2, dim=0),
+                        encoder_hidden_states=text_embeddings,
+                        cross_attention_kwargs={"scale": 0.0}
+                        if self.vsd_share_model
+                        else None,
+                        velocity_to_epsilon=self.pipe.scheduler.config.prediction_type
+                        == "v_prediction",
+                    )
+
+            noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+
+        return noise_pred
+
+    def get_eps_phi(
+        self,
+        latents_noisy: Float[Tensor, "B 4 Hl Wl"],
+        t: Int[Tensor, "B"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        camera_condition: Float[Tensor, "B ..."],
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        batch_size = latents_noisy.shape[0]
+
+        # not using view-dependent prompting in LoRA
+        text_embeddings, _ = prompt_utils.get_text_embeddings(
+            elevation, azimuth, camera_distances, view_dependent_prompting=False
+        ).chunk(2)
+        with torch.no_grad():
+            noise_pred = self.forward_unet(
+                self.pipe_phi.unet,
+                torch.cat([latents_noisy] * 2, dim=0),
+                torch.cat([t] * 2, dim=0),
+                encoder_hidden_states=torch.cat([text_embeddings] * 2, dim=0),
+                class_labels=torch.cat(
+                    [
+                        camera_condition.view(batch_size, -1),
+                        torch.zeros_like(camera_condition.view(batch_size, -1)),
+                    ],
+                    dim=0,
+                )
+                if self.cfg.vsd_use_camera_condition
+                else None,
+                cross_attention_kwargs={"scale": 1.0},
+                velocity_to_epsilon=self.pipe_phi.scheduler.config.prediction_type
+                == "v_prediction",
+            )
+
+        noise_pred_camera, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.vsd_guidance_scale_phi * (
+            noise_pred_camera - noise_pred_uncond
+        )
+
+        return noise_pred
+
+    def train_phi(
+        self,
+        latents: Float[Tensor, "B 4 Hl Wl"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        camera_condition: Float[Tensor, "B ..."],
+    ):
+        B = latents.shape[0]
+        latents = latents.detach().repeat(
+            self.cfg.vsd_lora_n_timestamp_samples, 1, 1, 1
+        )
+
+        num_train_timesteps = self.pipe_phi.scheduler.config.num_train_timesteps
+        t = torch.randint(
+            int(num_train_timesteps * 0.0),
+            int(num_train_timesteps * 1.0),
+            [B * self.cfg.vsd_lora_n_timestamp_samples],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        noise = torch.randn_like(latents)
+        latents_noisy = self.pipe_phi.scheduler.add_noise(latents, noise, t)
+        if self.pipe_phi.scheduler.config.prediction_type == "epsilon":
+            target = noise
+        elif self.pipe_phi.scheduler.prediction_type == "v_prediction":
+            target = self.pipe_phi.scheduler.get_velocity(latents, noise, t)
+        else:
+            raise ValueError(
+                f"Unknown prediction type {self.pipe_phi.scheduler.prediction_type}"
+            )
+
+        # not using view-dependent prompting in LoRA
+        text_embeddings, _ = prompt_utils.get_text_embeddings(
+            elevation, azimuth, camera_distances, view_dependent_prompting=False
+        ).chunk(2)
+
+        if (
+            self.cfg.vsd_use_camera_condition
+            and self.cfg.vsd_lora_cfg_training
+            and random.random() < 0.1
+        ):
+            camera_condition = torch.zeros_like(camera_condition)
+
+        noise_pred = self.forward_unet(
+            self.pipe_phi.unet,
+            latents_noisy,
+            t,
+            encoder_hidden_states=text_embeddings.repeat(
+                self.cfg.vsd_lora_n_timestamp_samples, 1, 1
+            ),
+            class_labels=camera_condition.view(B, -1).repeat(
+                self.cfg.vsd_lora_n_timestamp_samples, 1
+            )
+            if self.cfg.vsd_use_camera_condition
+            else None,
+            cross_attention_kwargs={"scale": 1.0},
+        )
+        return F.mse_loss(noise_pred.float(), target.float(), reduction="mean")
+
+    def forward(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        prompt_utils: PromptProcessorOutput,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        mvp_mtx: Float[Tensor, "B 4 4"],
+        c2w: Float[Tensor, "B 4 4"],
+        rgb_as_latents=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 Hl Wl"]
+        if rgb_as_latents:
+            # treat input rgb as latents
+            # input rgb should be in range [-1, 1]
+            latents = F.interpolate(
+                rgb_BCHW, (64, 64), mode="bilinear", align_corners=False
+            )
+        else:
+            # treat input rgb as rgb
+            # input rgb should be in range [0, 1]
+            rgb_BCHW = F.interpolate(
+                rgb_BCHW, (512, 512), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.vae_encode(self.pipe.vae, rgb_BCHW * 2.0 - 1.0)
+
+        # sample timestep
+        # use the same timestep for each batch
+        assert self.min_step is not None and self.max_step is not None
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [1],
+            dtype=torch.long,
+            device=self.device,
+        ).repeat(batch_size)
+
+        # sample noise
+        noise = torch.randn_like(latents)
+        latents_noisy = self.scheduler.add_noise(latents, noise, t)
+
+        eps_pretrain = self.get_eps_pretrain(
+            latents_noisy, t, prompt_utils, elevation, azimuth, camera_distances
+        )
+
+        latents_1step_orig = (
+            1
+            / self.alphas[t].view(-1, 1, 1, 1)
+            * (latents_noisy - self.sigmas[t].view(-1, 1, 1, 1) * eps_pretrain)
+        ).detach()
+
+        if self.cfg.guidance_type == "sds":
+            eps_phi = noise
+        elif self.cfg.guidance_type == "vsd":
+            if self.cfg.vsd_camera_condition_type == "extrinsics":
+                camera_condition = c2w
+            elif self.cfg.vsd_camera_condition_type == "mvp":
+                camera_condition = mvp_mtx
+            elif self.cfg.vsd_camera_condition_type == "spherical":
+                camera_condition = torch.stack(
+                    [
+                        torch.deg2rad(elevation),
+                        torch.sin(torch.deg2rad(azimuth)),
+                        torch.cos(torch.deg2rad(azimuth)),
+                        camera_distances,
+                    ],
+                    dim=-1,
+                )
+            else:
+                raise ValueError(
+                    f"Unknown camera_condition_type {self.cfg.vsd_camera_condition_type}"
+                )
+            eps_phi = self.get_eps_phi(
+                latents_noisy,
+                t,
+                prompt_utils,
+                elevation,
+                azimuth,
+                camera_distances,
+                camera_condition,
+            )
+
+            loss_train_phi = self.train_phi(
+                latents,
+                prompt_utils,
+                elevation,
+                azimuth,
+                camera_distances,
+                camera_condition,
+            )
+
+        if self.cfg.weighting_strategy == "dreamfusion":
+            w = (1.0 - self.alphas[t]).view(-1, 1, 1, 1)
+        elif self.cfg.weighting_strategy == "uniform":
+            w = 1.0
+        elif self.cfg.weighting_strategy == "fantasia3d":
+            w = (self.alphas[t] ** 0.5 * (1 - self.alphas[t])).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(
+                f"Unknown weighting strategy: {self.cfg.weighting_strategy}"
+            )
+
+        grad = w * (eps_pretrain - eps_phi)
+
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # reparameterization trick:
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        target = (latents - grad).detach()
+        loss_sd = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sd,
+            "grad_norm": grad.norm(),
+            "timesteps": t,
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+            "latents": latents,
+            "latents_1step_orig": latents_1step_orig,
+            "rgb": rgb_BCHW.permute(0, 2, 3, 1),
+            "weights": w,
+            "lambdas": self.lambdas[t],
+        }
+
+        if self.cfg.return_rgb_1step_orig:
+            with torch.no_grad():
+                rgb_1step_orig = self.vae_decode(
+                    self.pipe.vae, latents_1step_orig
+                ).permute(0, 2, 3, 1)
+            guidance_out.update({"rgb_1step_orig": rgb_1step_orig})
+
+        if self.cfg.return_rgb_multistep_orig:
+            with self.set_scheduler(
+                self.pipe,
+                DPMSolverSinglestepScheduler,
+                solver_order=1,
+                num_train_timesteps=int(t[0]),
+            ) as pipe:
+                text_embeddings = prompt_utils.get_text_embeddings(
+                    elevation,
+                    azimuth,
+                    camera_distances,
+                    self.cfg.view_dependent_prompting,
+                )
+                text_embeddings_cond, text_embeddings_uncond = text_embeddings.chunk(2)
+                with torch.cuda.amp.autocast(enabled=False):
+                    latents_multistep_orig = pipe(
+                        num_inference_steps=self.cfg.n_rgb_multistep_orig_steps,
+                        guidance_scale=self.cfg.guidance_scale,
+                        eta=1.0,
+                        latents=latents_noisy.to(pipe.unet.dtype),
+                        prompt_embeds=text_embeddings_cond.to(pipe.unet.dtype),
+                        negative_prompt_embeds=text_embeddings_uncond.to(
+                            pipe.unet.dtype
+                        ),
+                        cross_attention_kwargs={"scale": 0.0}
+                        if self.vsd_share_model
+                        else None,
+                        output_type="latent",
+                    ).images.to(latents.dtype)
+            with torch.no_grad():
+                rgb_multistep_orig = self.vae_decode(
+                    self.pipe.vae, latents_multistep_orig
+                )
+            guidance_out.update(
+                {
+                    "latents_multistep_orig": latents_multistep_orig,
+                    "rgb_multistep_orig": rgb_multistep_orig.permute(0, 2, 3, 1),
+                }
+            )
+
+        if self.cfg.guidance_type == "vsd":
+            guidance_out.update(
+                {
+                    "loss_train_phi": loss_train_phi,
+                }
+            )
+
+        return guidance_out
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.min_step = int(
+            self.num_train_timesteps * C(self.cfg.min_step_percent, epoch, global_step)
+        )
+        self.max_step = int(
+            self.num_train_timesteps * C(self.cfg.max_step_percent, epoch, global_step)
+        )

threestudio/models/guidance/stable_zero123_guidance.py

@@ -0,0 +1,340 @@
+import importlib
+import os
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import DDIMScheduler, DDPMScheduler, StableDiffusionPipeline
+from diffusers.utils.import_utils import is_xformers_available
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+import threestudio
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.typing import *
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        if config == "__is_first_stage__":
+            return None
+        elif config == "__is_unconditional__":
+            return None
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+# load model
+def load_model_from_config(config, ckpt, device, vram_O=True, verbose=False):
+    pl_sd = torch.load(ckpt, map_location="cpu")
+
+    if "global_step" in pl_sd and verbose:
+        print(f'[INFO] Global Step: {pl_sd["global_step"]}')
+
+    sd = pl_sd["state_dict"]
+
+    model = instantiate_from_config(config.model)
+    m, u = model.load_state_dict(sd, strict=False)
+
+    if len(m) > 0 and verbose:
+        print("[INFO] missing keys: \n", m)
+    if len(u) > 0 and verbose:
+        print("[INFO] unexpected keys: \n", u)
+
+    # manually load ema and delete it to save GPU memory
+    if model.use_ema:
+        if verbose:
+            print("[INFO] loading EMA...")
+        model.model_ema.copy_to(model.model)
+        del model.model_ema
+
+    if vram_O:
+        # we don't need decoder
+        del model.first_stage_model.decoder
+
+    torch.cuda.empty_cache()
+
+    model.eval().to(device)
+
+    return model
+
+
+@threestudio.register("stable-zero123-guidance")
+class StableZero123Guidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        pretrained_model_name_or_path: str = "load/zero123/stable-zero123.ckpt"
+        pretrained_config: str = "load/zero123/sd-objaverse-finetune-c_concat-256.yaml"
+        vram_O: bool = True
+
+        cond_image_path: str = "load/images/hamburger_rgba.png"
+        cond_elevation_deg: float = 0.0
+        cond_azimuth_deg: float = 0.0
+        cond_camera_distance: float = 1.2
+
+        guidance_scale: float = 5.0
+
+        grad_clip: Optional[
+            Any
+        ] = None  # field(default_factory=lambda: [0, 2.0, 8.0, 1000])
+        half_precision_weights: bool = False
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading Stable Zero123 ...")
+
+        self.config = OmegaConf.load(self.cfg.pretrained_config)
+        # TODO: seems it cannot load into fp16...
+        self.weights_dtype = torch.float32
+        self.model = load_model_from_config(
+            self.config,
+            self.cfg.pretrained_model_name_or_path,
+            device=self.device,
+            vram_O=self.cfg.vram_O,
+        )
+
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+
+        # timesteps: use diffuser for convenience... hope it's alright.
+        self.num_train_timesteps = self.config.model.params.timesteps
+
+        self.scheduler = DDIMScheduler(
+            self.num_train_timesteps,
+            self.config.model.params.linear_start,
+            self.config.model.params.linear_end,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            set_alpha_to_one=False,
+            steps_offset=1,
+        )
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.grad_clip_val: Optional[float] = None
+
+        self.prepare_embeddings(self.cfg.cond_image_path)
+
+        threestudio.info(f"Loaded Stable Zero123!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def prepare_embeddings(self, image_path: str) -> None:
+        # load cond image for zero123
+        assert os.path.exists(image_path)
+        rgba = cv2.cvtColor(
+            cv2.imread(image_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA
+        )
+        rgba = (
+            cv2.resize(rgba, (256, 256), interpolation=cv2.INTER_AREA).astype(
+                np.float32
+            )
+            / 255.0
+        )
+        rgb = rgba[..., :3] * rgba[..., 3:] + (1 - rgba[..., 3:])
+        self.rgb_256: Float[Tensor, "1 3 H W"] = (
+            torch.from_numpy(rgb)
+            .unsqueeze(0)
+            .permute(0, 3, 1, 2)
+            .contiguous()
+            .to(self.device)
+        )
+        self.c_crossattn, self.c_concat = self.get_img_embeds(self.rgb_256)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_img_embeds(
+        self,
+        img: Float[Tensor, "B 3 256 256"],
+    ) -> Tuple[Float[Tensor, "B 1 768"], Float[Tensor, "B 4 32 32"]]:
+        img = img * 2.0 - 1.0
+        c_crossattn = self.model.get_learned_conditioning(img.to(self.weights_dtype))
+        c_concat = self.model.encode_first_stage(img.to(self.weights_dtype)).mode()
+        return c_crossattn, c_concat
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 256 256"]
+    ) -> Float[Tensor, "B 4 32 32"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        latents = self.model.get_first_stage_encoding(
+            self.model.encode_first_stage(imgs.to(self.weights_dtype))
+        )
+        return latents.to(input_dtype)  # [B, 4, 32, 32] Latent space image
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+    ) -> Float[Tensor, "B 3 512 512"]:
+        input_dtype = latents.dtype
+        image = self.model.decode_first_stage(latents)
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_cond(
+        self,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        c_crossattn=None,
+        c_concat=None,
+        **kwargs,
+    ) -> dict:
+        T = torch.stack(
+            [
+                torch.deg2rad(
+                    (90 - elevation) - (90 - self.cfg.cond_elevation_deg)
+                ),  # Zero123 polar is 90-elevation
+                torch.sin(torch.deg2rad(azimuth - self.cfg.cond_azimuth_deg)),
+                torch.cos(torch.deg2rad(azimuth - self.cfg.cond_azimuth_deg)),
+                torch.deg2rad(
+                    90 - torch.full_like(elevation, self.cfg.cond_elevation_deg)
+                ),
+            ],
+            dim=-1,
+        )[:, None, :].to(self.device)
+        cond = {}
+        clip_emb = self.model.cc_projection(
+            torch.cat(
+                [
+                    (self.c_crossattn if c_crossattn is None else c_crossattn).repeat(
+                        len(T), 1, 1
+                    ),
+                    T,
+                ],
+                dim=-1,
+            )
+        )
+        cond["c_crossattn"] = [
+            torch.cat([torch.zeros_like(clip_emb).to(self.device), clip_emb], dim=0)
+        ]
+        cond["c_concat"] = [
+            torch.cat(
+                [
+                    torch.zeros_like(self.c_concat)
+                    .repeat(len(T), 1, 1, 1)
+                    .to(self.device),
+                    (self.c_concat if c_concat is None else c_concat).repeat(
+                        len(T), 1, 1, 1
+                    ),
+                ],
+                dim=0,
+            )
+        ]
+        return cond
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        rgb_as_latents=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 64 64"]
+        if rgb_as_latents:
+            latents = (
+                F.interpolate(rgb_BCHW, (32, 32), mode="bilinear", align_corners=False)
+                * 2
+                - 1
+            )
+        else:
+            rgb_BCHW_512 = F.interpolate(
+                rgb_BCHW, (256, 256), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.encode_images(rgb_BCHW_512)
+
+        cond = self.get_cond(elevation, azimuth, camera_distances)
+
+        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [batch_size],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        # predict the noise residual with unet, NO grad!
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)  # TODO: use torch generator
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            x_in = torch.cat([latents_noisy] * 2)
+            t_in = torch.cat([t] * 2)
+            noise_pred = self.model.apply_model(x_in, t_in, cond)
+
+        # perform guidance
+        noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_cond - noise_pred_uncond
+        )
+
+        w = (1 - self.alphas[t]).reshape(-1, 1, 1, 1)
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+        # clip grad for stable training?
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # loss = SpecifyGradient.apply(latents, grad)
+        # SpecifyGradient is not straghtforward, use a reparameterization trick instead
+        target = (latents - grad).detach()
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sds,
+            "grad_norm": grad.norm(),
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+        }
+
+        return guidance_out
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )

threestudio/models/guidance/zero123_guidance.py

@@ -0,0 +1,528 @@
+import importlib
+import os
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import DDIMScheduler, DDPMScheduler, StableDiffusionPipeline
+from diffusers.utils.import_utils import is_xformers_available
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+import threestudio
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import C, parse_version
+from threestudio.utils.typing import *
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        if config == "__is_first_stage__":
+            return None
+        elif config == "__is_unconditional__":
+            return None
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+# load model
+def load_model_from_config(config, ckpt, device, vram_O=True, verbose=False):
+    pl_sd = torch.load(ckpt, map_location="cpu")
+
+    if "global_step" in pl_sd and verbose:
+        print(f'[INFO] Global Step: {pl_sd["global_step"]}')
+
+    sd = pl_sd["state_dict"]
+
+    model = instantiate_from_config(config.model)
+    m, u = model.load_state_dict(sd, strict=False)
+
+    if len(m) > 0 and verbose:
+        print("[INFO] missing keys: \n", m)
+    if len(u) > 0 and verbose:
+        print("[INFO] unexpected keys: \n", u)
+
+    # manually load ema and delete it to save GPU memory
+    if model.use_ema:
+        if verbose:
+            print("[INFO] loading EMA...")
+        model.model_ema.copy_to(model.model)
+        del model.model_ema
+
+    if vram_O:
+        # we don't need decoder
+        del model.first_stage_model.decoder
+
+    torch.cuda.empty_cache()
+
+    model.eval().to(device)
+
+    return model
+
+
+@threestudio.register("zero123-guidance")
+class Zero123Guidance(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        pretrained_model_name_or_path: str = "load/zero123/105000.ckpt"
+        pretrained_config: str = "load/zero123/sd-objaverse-finetune-c_concat-256.yaml"
+        vram_O: bool = True
+
+        cond_image_path: str = "load/images/hamburger_rgba.png"
+        cond_elevation_deg: float = 0.0
+        cond_azimuth_deg: float = 0.0
+        cond_camera_distance: float = 1.2
+
+        guidance_scale: float = 5.0
+
+        grad_clip: Optional[
+            Any
+        ] = None  # field(default_factory=lambda: [0, 2.0, 8.0, 1000])
+        half_precision_weights: bool = False
+
+        min_step_percent: float = 0.02
+        max_step_percent: float = 0.98
+
+        """Maximum number of batch items to evaluate guidance for (for debugging) and to save on disk. -1 means save all items."""
+        max_items_eval: int = 4
+
+    cfg: Config
+
+    def configure(self) -> None:
+        threestudio.info(f"Loading Zero123 ...")
+
+        self.config = OmegaConf.load(self.cfg.pretrained_config)
+        # TODO: seems it cannot load into fp16...
+        self.weights_dtype = torch.float32
+        self.model = load_model_from_config(
+            self.config,
+            self.cfg.pretrained_model_name_or_path,
+            device=self.device,
+            vram_O=self.cfg.vram_O,
+        )
+
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+
+        # timesteps: use diffuser for convenience... hope it's alright.
+        self.num_train_timesteps = self.config.model.params.timesteps
+
+        self.scheduler = DDIMScheduler(
+            self.num_train_timesteps,
+            self.config.model.params.linear_start,
+            self.config.model.params.linear_end,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            set_alpha_to_one=False,
+            steps_offset=1,
+        )
+
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+        self.set_min_max_steps()  # set to default value
+
+        self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+
+        self.grad_clip_val: Optional[float] = None
+
+        self.prepare_embeddings(self.cfg.cond_image_path)
+
+        threestudio.info(f"Loaded Zero123!")
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def set_min_max_steps(self, min_step_percent=0.02, max_step_percent=0.98):
+        self.min_step = int(self.num_train_timesteps * min_step_percent)
+        self.max_step = int(self.num_train_timesteps * max_step_percent)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def prepare_embeddings(self, image_path: str) -> None:
+        # load cond image for zero123
+        assert os.path.exists(image_path)
+        rgba = cv2.cvtColor(
+            cv2.imread(image_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA
+        )
+        rgba = (
+            cv2.resize(rgba, (256, 256), interpolation=cv2.INTER_AREA).astype(
+                np.float32
+            )
+            / 255.0
+        )
+        rgb = rgba[..., :3] * rgba[..., 3:] + (1 - rgba[..., 3:])
+        self.rgb_256: Float[Tensor, "1 3 H W"] = (
+            torch.from_numpy(rgb)
+            .unsqueeze(0)
+            .permute(0, 3, 1, 2)
+            .contiguous()
+            .to(self.device)
+        )
+        self.c_crossattn, self.c_concat = self.get_img_embeds(self.rgb_256)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_img_embeds(
+        self,
+        img: Float[Tensor, "B 3 256 256"],
+    ) -> Tuple[Float[Tensor, "B 1 768"], Float[Tensor, "B 4 32 32"]]:
+        img = img * 2.0 - 1.0
+        c_crossattn = self.model.get_learned_conditioning(img.to(self.weights_dtype))
+        c_concat = self.model.encode_first_stage(img.to(self.weights_dtype)).mode()
+        return c_crossattn, c_concat
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def encode_images(
+        self, imgs: Float[Tensor, "B 3 256 256"]
+    ) -> Float[Tensor, "B 4 32 32"]:
+        input_dtype = imgs.dtype
+        imgs = imgs * 2.0 - 1.0
+        latents = self.model.get_first_stage_encoding(
+            self.model.encode_first_stage(imgs.to(self.weights_dtype))
+        )
+        return latents.to(input_dtype)  # [B, 4, 32, 32] Latent space image
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def decode_latents(
+        self,
+        latents: Float[Tensor, "B 4 H W"],
+    ) -> Float[Tensor, "B 3 512 512"]:
+        input_dtype = latents.dtype
+        image = self.model.decode_first_stage(latents)
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def get_cond(
+        self,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        c_crossattn=None,
+        c_concat=None,
+        **kwargs,
+    ) -> dict:
+        T = torch.stack(
+            [
+                torch.deg2rad(
+                    (90 - elevation) - (90 - self.cfg.cond_elevation_deg)
+                ),  # Zero123 polar is 90-elevation
+                torch.sin(torch.deg2rad(azimuth - self.cfg.cond_azimuth_deg)),
+                torch.cos(torch.deg2rad(azimuth - self.cfg.cond_azimuth_deg)),
+                camera_distances - self.cfg.cond_camera_distance,
+            ],
+            dim=-1,
+        )[:, None, :].to(self.device)
+        cond = {}
+        clip_emb = self.model.cc_projection(
+            torch.cat(
+                [
+                    (self.c_crossattn if c_crossattn is None else c_crossattn).repeat(
+                        len(T), 1, 1
+                    ),
+                    T,
+                ],
+                dim=-1,
+            )
+        )
+        cond["c_crossattn"] = [
+            torch.cat([torch.zeros_like(clip_emb).to(self.device), clip_emb], dim=0)
+        ]
+        cond["c_concat"] = [
+            torch.cat(
+                [
+                    torch.zeros_like(self.c_concat)
+                    .repeat(len(T), 1, 1, 1)
+                    .to(self.device),
+                    (self.c_concat if c_concat is None else c_concat).repeat(
+                        len(T), 1, 1, 1
+                    ),
+                ],
+                dim=0,
+            )
+        ]
+        return cond
+
+    def __call__(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        rgb_as_latents=False,
+        guidance_eval=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 64 64"]
+        if rgb_as_latents:
+            latents = (
+                F.interpolate(rgb_BCHW, (32, 32), mode="bilinear", align_corners=False)
+                * 2
+                - 1
+            )
+        else:
+            rgb_BCHW_512 = F.interpolate(
+                rgb_BCHW, (256, 256), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.encode_images(rgb_BCHW_512)
+
+        cond = self.get_cond(elevation, azimuth, camera_distances)
+
+        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [batch_size],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        # predict the noise residual with unet, NO grad!
+        with torch.no_grad():
+            # add noise
+            noise = torch.randn_like(latents)  # TODO: use torch generator
+            latents_noisy = self.scheduler.add_noise(latents, noise, t)
+            # pred noise
+            x_in = torch.cat([latents_noisy] * 2)
+            t_in = torch.cat([t] * 2)
+            noise_pred = self.model.apply_model(x_in, t_in, cond)
+
+        # perform guidance
+        noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_cond - noise_pred_uncond
+        )
+
+        w = (1 - self.alphas[t]).reshape(-1, 1, 1, 1)
+        grad = w * (noise_pred - noise)
+        grad = torch.nan_to_num(grad)
+        # clip grad for stable training?
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # loss = SpecifyGradient.apply(latents, grad)
+        # SpecifyGradient is not straghtforward, use a reparameterization trick instead
+        target = (latents - grad).detach()
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        loss_sds = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sds, # loss_sds
+            "grad_norm": grad.norm(),
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+        }
+
+        if guidance_eval:
+            guidance_eval_utils = {
+                "cond": cond,
+                "t_orig": t,
+                "latents_noisy": latents_noisy,
+                "noise_pred": noise_pred,
+            }
+            guidance_eval_out = self.guidance_eval(**guidance_eval_utils)
+            texts = []
+            for n, e, a, c in zip(
+                guidance_eval_out["noise_levels"], elevation, azimuth, camera_distances
+            ):
+                texts.append(
+                    f"n{n:.02f}\ne{e.item():.01f}\na{a.item():.01f}\nc{c.item():.02f}"
+                )
+            guidance_eval_out.update({"texts": texts})
+            guidance_out.update({"eval": guidance_eval_out})
+
+        return guidance_out
+
+    @torch.cuda.amp.autocast(enabled=False)
+    @torch.no_grad()
+    def guidance_eval(self, cond, t_orig, latents_noisy, noise_pred):
+        # use only 50 timesteps, and find nearest of those to t
+        self.scheduler.set_timesteps(50)
+        self.scheduler.timesteps_gpu = self.scheduler.timesteps.to(self.device)
+        bs = (
+            min(self.cfg.max_items_eval, latents_noisy.shape[0])
+            if self.cfg.max_items_eval > 0
+            else latents_noisy.shape[0]
+        )  # batch size
+        large_enough_idxs = self.scheduler.timesteps_gpu.expand([bs, -1]) > t_orig[
+            :bs
+        ].unsqueeze(
+            -1
+        )  # sized [bs,50] > [bs,1]
+        idxs = torch.min(large_enough_idxs, dim=1)[1]
+        t = self.scheduler.timesteps_gpu[idxs]
+
+        fracs = list((t / self.scheduler.config.num_train_timesteps).cpu().numpy())
+        imgs_noisy = self.decode_latents(latents_noisy[:bs]).permute(0, 2, 3, 1)
+
+        # get prev latent
+        latents_1step = []
+        pred_1orig = []
+        for b in range(bs):
+            step_output = self.scheduler.step(
+                noise_pred[b : b + 1], t[b], latents_noisy[b : b + 1], eta=1
+            )
+            latents_1step.append(step_output["prev_sample"])
+            pred_1orig.append(step_output["pred_original_sample"])
+        latents_1step = torch.cat(latents_1step)
+        pred_1orig = torch.cat(pred_1orig)
+        imgs_1step = self.decode_latents(latents_1step).permute(0, 2, 3, 1)
+        imgs_1orig = self.decode_latents(pred_1orig).permute(0, 2, 3, 1)
+
+        latents_final = []
+        for b, i in enumerate(idxs):
+            latents = latents_1step[b : b + 1]
+            c = {
+                "c_crossattn": [cond["c_crossattn"][0][[b, b + len(idxs)], ...]],
+                "c_concat": [cond["c_concat"][0][[b, b + len(idxs)], ...]],
+            }
+            for t in tqdm(self.scheduler.timesteps[i + 1 :], leave=False):
+                # pred noise
+                x_in = torch.cat([latents] * 2)
+                t_in = torch.cat([t.reshape(1)] * 2).to(self.device)
+                noise_pred = self.model.apply_model(x_in, t_in, c)
+                # perform guidance
+                noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+                    noise_pred_cond - noise_pred_uncond
+                )
+                # get prev latent
+                latents = self.scheduler.step(noise_pred, t, latents, eta=1)[
+                    "prev_sample"
+                ]
+            latents_final.append(latents)
+
+        latents_final = torch.cat(latents_final)
+        imgs_final = self.decode_latents(latents_final).permute(0, 2, 3, 1)
+
+        return {
+            "bs": bs,
+            "noise_levels": fracs,
+            "imgs_noisy": imgs_noisy,
+            "imgs_1step": imgs_1step,
+            "imgs_1orig": imgs_1orig,
+            "imgs_final": imgs_final,
+        }
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.set_min_max_steps(
+            min_step_percent=C(self.cfg.min_step_percent, epoch, global_step),
+            max_step_percent=C(self.cfg.max_step_percent, epoch, global_step),
+        )
+
+    # verification - requires `vram_O = False` in load_model_from_config
+    @torch.no_grad()
+    def generate(
+        self,
+        image,  # image tensor [1, 3, H, W] in [0, 1]
+        elevation=0,
+        azimuth=0,
+        camera_distances=0,  # new view params
+        c_crossattn=None,
+        c_concat=None,
+        scale=3,
+        ddim_steps=50,
+        post_process=True,
+        ddim_eta=1,
+    ):
+        if c_crossattn is None:
+            c_crossattn, c_concat = self.get_img_embeds(image)
+
+        cond = self.get_cond(
+            elevation, azimuth, camera_distances, c_crossattn, c_concat
+        )
+
+        imgs = self.gen_from_cond(cond, scale, ddim_steps, post_process, ddim_eta)
+
+        return imgs
+
+    # verification - requires `vram_O = False` in load_model_from_config
+    @torch.no_grad()
+    def gen_from_cond(
+        self,
+        cond,
+        scale=3,
+        ddim_steps=50,
+        post_process=True,
+        ddim_eta=1,
+    ):
+        # produce latents loop
+        B = cond["c_crossattn"][0].shape[0] // 2
+        latents = torch.randn((B, 4, 32, 32), device=self.device)
+        self.scheduler.set_timesteps(ddim_steps)
+
+        for t in self.scheduler.timesteps:
+            x_in = torch.cat([latents] * 2)
+            t_in = torch.cat([t.reshape(1).repeat(B)] * 2).to(self.device)
+
+            noise_pred = self.model.apply_model(x_in, t_in, cond)
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + scale * (
+                noise_pred_cond - noise_pred_uncond
+            )
+
+            latents = self.scheduler.step(noise_pred, t, latents, eta=ddim_eta)[
+                "prev_sample"
+            ]
+
+        imgs = self.decode_latents(latents)
+        imgs = imgs.cpu().numpy().transpose(0, 2, 3, 1) if post_process else imgs
+
+        return imgs
+
+
+if __name__ == '__main__':
+    from threestudio.utils.config import load_config
+    import pytorch_lightning as pl
+    import numpy as np
+    import os
+    import cv2
+    cfg = load_config("configs/experimental/zero123.yaml")
+    guidance = threestudio.find(cfg.system.guidance_type)(cfg.system.guidance)
+    elevations = [0, 20, -20]
+    azimuths = [45, 90, 135, -45, -90]
+    radius = torch.tensor([3.8]).to(guidance.device)
+    outdir = ".threestudio_cache/saiyan"
+    os.makedirs(outdir, exist_ok=True)
+    # rgb_image = (rgb_image[0].detach().cpu().clip(0, 1).numpy()*255).astype(np.uint8)[:, :, ::-1].copy()
+    # os.makedirs('.threestudio_cache', exist_ok=True)
+    # cv2.imwrite('.threestudio_cache/diffusion_image.jpg', rgb_image)
+
+
+    rgb_image = cv2.imread(cfg.system.guidance.cond_image_path)[:, :, ::-1].copy() / 255
+    rgb_image = cv2.resize(rgb_image, (256, 256))
+    rgb_image = torch.FloatTensor(rgb_image).unsqueeze(0).to(guidance.device).permute(0,3,1,2)
+
+    for elevation in elevations:
+        for azimuth in azimuths:
+            output1 = guidance.generate(
+                rgb_image,
+                torch.tensor([elevation]).to(guidance.device),
+                torch.tensor([azimuth]).to(guidance.device),
+                radius,
+                c_crossattn=guidance.c_crossattn,
+                c_concat=guidance.c_concat
+            )
+            from torchvision.utils import save_image
+            save_image(torch.tensor(output1).float().permute(0,3,1,2), f"{outdir}/result_e_{elevation}_a_{azimuth}.png", normalize=True, value_range=(0,1))
+            

threestudio/models/guidance/zero123_unified_guidance.py

@@ -0,0 +1,721 @@
+import os
+import random
+import sys
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    DPMSolverSinglestepScheduler,
+    UNet2DConditionModel,
+)
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import LoRAAttnProcessor
+from diffusers.models.embeddings import TimestepEmbedding
+from PIL import Image
+from tqdm import tqdm
+
+import threestudio
+from extern.zero123 import Zero123Pipeline
+from threestudio.models.networks import ToDTypeWrapper
+from threestudio.models.prompt_processors.base import PromptProcessorOutput
+from threestudio.utils.base import BaseModule
+from threestudio.utils.misc import C, cleanup, enable_gradient, parse_version
+from threestudio.utils.typing import *
+
+
+@threestudio.register("zero123-unified-guidance")
+class Zero123UnifiedGuidance(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        cache_dir: Optional[str] = None
+        local_files_only: Optional[bool] = False
+
+        # guidance type, in ["sds", "vsd"]
+        guidance_type: str = "sds"
+
+        pretrained_model_name_or_path: str = "bennyguo/zero123-diffusers"
+        guidance_scale: float = 5.0
+        weighting_strategy: str = "dreamfusion"
+
+        min_step_percent: Any = 0.02
+        max_step_percent: Any = 0.98
+        grad_clip: Optional[Any] = None
+
+        return_rgb_1step_orig: bool = False
+        return_rgb_multistep_orig: bool = False
+        n_rgb_multistep_orig_steps: int = 4
+
+        cond_image_path: str = ""
+        cond_elevation_deg: float = 0.0
+        cond_azimuth_deg: float = 0.0
+        cond_camera_distance: float = 1.2
+
+        # efficiency-related configurations
+        half_precision_weights: bool = True
+
+        # VSD configurations, only used when guidance_type is "vsd"
+        vsd_phi_model_name_or_path: Optional[str] = None
+        vsd_guidance_scale_phi: float = 1.0
+        vsd_use_lora: bool = True
+        vsd_lora_cfg_training: bool = False
+        vsd_lora_n_timestamp_samples: int = 1
+        vsd_use_camera_condition: bool = True
+        # camera condition type, in ["extrinsics", "mvp", "spherical"]
+        vsd_camera_condition_type: Optional[str] = "extrinsics"
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.min_step: Optional[int] = None
+        self.max_step: Optional[int] = None
+        self.grad_clip_val: Optional[float] = None
+
+        @dataclass
+        class NonTrainableModules:
+            pipe: Zero123Pipeline
+            pipe_phi: Optional[Zero123Pipeline] = None
+
+        self.weights_dtype = (
+            torch.float16 if self.cfg.half_precision_weights else torch.float32
+        )
+
+        threestudio.info(f"Loading Zero123 ...")
+
+        # need to make sure the pipeline file is in path
+        sys.path.append("extern/")
+
+        pipe_kwargs = {
+            "safety_checker": None,
+            "requires_safety_checker": False,
+            "variant": "fp16" if self.cfg.half_precision_weights else None,
+            "torch_dtype": self.weights_dtype,
+            "cache_dir": self.cfg.cache_dir,
+            "local_files_only": self.cfg.local_files_only,
+        }
+        pipe = Zero123Pipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            **pipe_kwargs,
+        ).to(self.device)
+        self.prepare_pipe(pipe)
+
+        # phi network for VSD
+        # introduce two trainable modules:
+        # - self.camera_embedding
+        # - self.lora_layers
+        pipe_phi = None
+
+        # if the phi network shares the same unet with the pretrain network
+        # we need to pass additional cross attention kwargs to the unet
+        self.vsd_share_model = (
+            self.cfg.guidance_type == "vsd"
+            and self.cfg.vsd_phi_model_name_or_path is None
+        )
+        if self.cfg.guidance_type == "vsd":
+            if self.cfg.vsd_phi_model_name_or_path is None:
+                pipe_phi = pipe
+            else:
+                pipe_phi = Zero123Pipeline.from_pretrained(
+                    self.cfg.vsd_phi_model_name_or_path,
+                    **pipe_kwargs,
+                ).to(self.device)
+                self.prepare_pipe(pipe_phi)
+
+            # set up camera embedding
+            if self.cfg.vsd_use_camera_condition:
+                if self.cfg.vsd_camera_condition_type in ["extrinsics", "mvp"]:
+                    self.camera_embedding_dim = 16
+                elif self.cfg.vsd_camera_condition_type == "spherical":
+                    self.camera_embedding_dim = 4
+                else:
+                    raise ValueError("Invalid camera condition type!")
+
+                # FIXME: hard-coded output dim
+                self.camera_embedding = ToDTypeWrapper(
+                    TimestepEmbedding(self.camera_embedding_dim, 1280),
+                    self.weights_dtype,
+                ).to(self.device)
+                pipe_phi.unet.class_embedding = self.camera_embedding
+
+            if self.cfg.vsd_use_lora:
+                # set up LoRA layers
+                lora_attn_procs = {}
+                for name in pipe_phi.unet.attn_processors.keys():
+                    cross_attention_dim = (
+                        None
+                        if name.endswith("attn1.processor")
+                        else pipe_phi.unet.config.cross_attention_dim
+                    )
+                    if name.startswith("mid_block"):
+                        hidden_size = pipe_phi.unet.config.block_out_channels[-1]
+                    elif name.startswith("up_blocks"):
+                        block_id = int(name[len("up_blocks.")])
+                        hidden_size = list(
+                            reversed(pipe_phi.unet.config.block_out_channels)
+                        )[block_id]
+                    elif name.startswith("down_blocks"):
+                        block_id = int(name[len("down_blocks.")])
+                        hidden_size = pipe_phi.unet.config.block_out_channels[block_id]
+
+                    lora_attn_procs[name] = LoRAAttnProcessor(
+                        hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
+                    )
+
+                pipe_phi.unet.set_attn_processor(lora_attn_procs)
+
+                self.lora_layers = AttnProcsLayers(pipe_phi.unet.attn_processors).to(
+                    self.device
+                )
+                self.lora_layers._load_state_dict_pre_hooks.clear()
+                self.lora_layers._state_dict_hooks.clear()
+
+        threestudio.info(f"Loaded Stable Diffusion!")
+
+        self.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+        self.num_train_timesteps = self.scheduler.config.num_train_timesteps
+
+        # q(z_t|x) = N(alpha_t x, sigma_t^2 I)
+        # in DDPM, alpha_t = sqrt(alphas_cumprod_t), sigma_t^2 = 1 - alphas_cumprod_t
+        self.alphas_cumprod: Float[Tensor, "T"] = self.scheduler.alphas_cumprod.to(
+            self.device
+        )
+        self.alphas: Float[Tensor, "T"] = self.alphas_cumprod**0.5
+        self.sigmas: Float[Tensor, "T"] = (1 - self.alphas_cumprod) ** 0.5
+        # log SNR
+        self.lambdas: Float[Tensor, "T"] = self.sigmas / self.alphas
+
+        self._non_trainable_modules = NonTrainableModules(
+            pipe=pipe,
+            pipe_phi=pipe_phi,
+        )
+
+        # self.clip_image_embeddings and self.image_latents
+        self.prepare_image_embeddings()
+
+    @property
+    def pipe(self) -> Zero123Pipeline:
+        return self._non_trainable_modules.pipe
+
+    @property
+    def pipe_phi(self) -> Zero123Pipeline:
+        if self._non_trainable_modules.pipe_phi is None:
+            raise RuntimeError("phi model is not available.")
+        return self._non_trainable_modules.pipe_phi
+
+    def prepare_pipe(self, pipe: Zero123Pipeline):
+        cleanup()
+
+        pipe.image_encoder.eval()
+        pipe.vae.eval()
+        pipe.unet.eval()
+        pipe.clip_camera_projection.eval()
+
+        enable_gradient(pipe.image_encoder, enabled=False)
+        enable_gradient(pipe.vae, enabled=False)
+        enable_gradient(pipe.unet, enabled=False)
+        enable_gradient(pipe.clip_camera_projection, enabled=False)
+
+        # disable progress bar
+        pipe.set_progress_bar_config(disable=True)
+
+    def prepare_image_embeddings(self) -> None:
+        if not os.path.exists(self.cfg.cond_image_path):
+            raise RuntimeError(
+                f"Condition image not found at {self.cfg.cond_image_path}"
+            )
+        image = Image.open(self.cfg.cond_image_path).convert("RGBA").resize((256, 256))
+        image = (
+            TF.to_tensor(image)
+            .unsqueeze(0)
+            .to(device=self.device, dtype=self.weights_dtype)
+        )
+        # rgba -> rgb, apply white background
+        image = image[:, :3] * image[:, 3:4] + (1 - image[:, 3:4])
+
+        with torch.no_grad():
+            self.clip_image_embeddings: Float[
+                Tensor, "1 1 D"
+            ] = self.extract_clip_image_embeddings(image)
+
+            # encoded latents should be multiplied with vae.config.scaling_factor
+            # but zero123 was not trained this way
+            self.image_latents: Float[Tensor, "1 4 Hl Wl"] = (
+                self.vae_encode(self.pipe.vae, image * 2.0 - 1.0, mode=True)
+                / self.pipe.vae.config.scaling_factor
+            )
+
+    def extract_clip_image_embeddings(
+        self, images: Float[Tensor, "B 3 H W"]
+    ) -> Float[Tensor, "B 1 D"]:
+        # expect images in [0, 1]
+        images_pil = [TF.to_pil_image(image) for image in images]
+        images_processed = self.pipe.feature_extractor(
+            images=images_pil, return_tensors="pt"
+        ).pixel_values.to(device=self.device, dtype=self.weights_dtype)
+        clip_image_embeddings = self.pipe.image_encoder(images_processed).image_embeds
+        return clip_image_embeddings.to(images.dtype)
+
+    def get_image_camera_embeddings(
+        self,
+        elevation_deg: Float[Tensor, "B"],
+        azimuth_deg: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ) -> Float[Tensor, "B 1 D"]:
+        batch_size = elevation_deg.shape[0]
+        camera_embeddings: Float[Tensor, "B 1 4"] = torch.stack(
+            [
+                torch.deg2rad(self.cfg.cond_elevation_deg - elevation_deg),
+                torch.sin(torch.deg2rad(azimuth_deg - self.cfg.cond_azimuth_deg)),
+                torch.cos(torch.deg2rad(azimuth_deg - self.cfg.cond_azimuth_deg)),
+                camera_distances - self.cfg.cond_camera_distance,
+            ],
+            dim=-1,
+        )[:, None, :]
+
+        image_camera_embeddings = self.pipe.clip_camera_projection(
+            torch.cat(
+                [
+                    self.clip_image_embeddings.repeat(batch_size, 1, 1),
+                    camera_embeddings,
+                ],
+                dim=-1,
+            ).to(self.weights_dtype)
+        )
+
+        return image_camera_embeddings
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward_unet(
+        self,
+        unet: UNet2DConditionModel,
+        latents: Float[Tensor, "..."],
+        t: Int[Tensor, "..."],
+        encoder_hidden_states: Float[Tensor, "..."],
+        class_labels: Optional[Float[Tensor, "..."]] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        down_block_additional_residuals: Optional[Float[Tensor, "..."]] = None,
+        mid_block_additional_residual: Optional[Float[Tensor, "..."]] = None,
+        velocity_to_epsilon: bool = False,
+    ) -> Float[Tensor, "..."]:
+        input_dtype = latents.dtype
+        pred = unet(
+            latents.to(unet.dtype),
+            t.to(unet.dtype),
+            encoder_hidden_states=encoder_hidden_states.to(unet.dtype),
+            class_labels=class_labels,
+            cross_attention_kwargs=cross_attention_kwargs,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+        ).sample
+        if velocity_to_epsilon:
+            pred = latents * self.sigmas[t].view(-1, 1, 1, 1) + pred * self.alphas[
+                t
+            ].view(-1, 1, 1, 1)
+        return pred.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def vae_encode(
+        self, vae: AutoencoderKL, imgs: Float[Tensor, "B 3 H W"], mode=False
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        # expect input in [-1, 1]
+        input_dtype = imgs.dtype
+        posterior = vae.encode(imgs.to(vae.dtype)).latent_dist
+        if mode:
+            latents = posterior.mode()
+        else:
+            latents = posterior.sample()
+        latents = latents * vae.config.scaling_factor
+        return latents.to(input_dtype)
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def vae_decode(
+        self, vae: AutoencoderKL, latents: Float[Tensor, "B 4 Hl Wl"]
+    ) -> Float[Tensor, "B 3 H W"]:
+        # output in [0, 1]
+        input_dtype = latents.dtype
+        latents = 1 / vae.config.scaling_factor * latents
+        image = vae.decode(latents.to(vae.dtype)).sample
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        return image.to(input_dtype)
+
+    @contextmanager
+    def disable_unet_class_embedding(self, unet: UNet2DConditionModel):
+        class_embedding = unet.class_embedding
+        try:
+            unet.class_embedding = None
+            yield unet
+        finally:
+            unet.class_embedding = class_embedding
+
+    @contextmanager
+    def set_scheduler(self, pipe: Zero123Pipeline, scheduler_class: Any, **kwargs):
+        scheduler_orig = pipe.scheduler
+        pipe.scheduler = scheduler_class.from_config(scheduler_orig.config, **kwargs)
+        yield pipe
+        pipe.scheduler = scheduler_orig
+
+    def get_eps_pretrain(
+        self,
+        latents_noisy: Float[Tensor, "B 4 Hl Wl"],
+        t: Int[Tensor, "B"],
+        image_camera_embeddings: Float[Tensor, "B 1 D"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        batch_size = latents_noisy.shape[0]
+
+        with torch.no_grad():
+            with self.disable_unet_class_embedding(self.pipe.unet) as unet:
+                noise_pred = self.forward_unet(
+                    unet,
+                    torch.cat(
+                        [
+                            torch.cat([latents_noisy] * 2, dim=0),
+                            torch.cat(
+                                [
+                                    self.image_latents.repeat(batch_size, 1, 1, 1),
+                                    torch.zeros_like(self.image_latents).repeat(
+                                        batch_size, 1, 1, 1
+                                    ),
+                                ],
+                                dim=0,
+                            ),
+                        ],
+                        dim=1,
+                    ),
+                    torch.cat([t] * 2, dim=0),
+                    encoder_hidden_states=torch.cat(
+                        [
+                            image_camera_embeddings,
+                            torch.zeros_like(image_camera_embeddings),
+                        ],
+                        dim=0,
+                    ),
+                    cross_attention_kwargs={"scale": 0.0}
+                    if self.vsd_share_model
+                    else None,
+                    velocity_to_epsilon=self.pipe.scheduler.config.prediction_type
+                    == "v_prediction",
+                )
+
+        noise_pred_image, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.guidance_scale * (
+            noise_pred_image - noise_pred_uncond
+        )
+
+        return noise_pred
+
+    def get_eps_phi(
+        self,
+        latents_noisy: Float[Tensor, "B 4 Hl Wl"],
+        t: Int[Tensor, "B"],
+        image_camera_embeddings: Float[Tensor, "B 1 D"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        camera_condition: Float[Tensor, "B ..."],
+    ) -> Float[Tensor, "B 4 Hl Wl"]:
+        batch_size = latents_noisy.shape[0]
+
+        with torch.no_grad():
+            noise_pred = self.forward_unet(
+                self.pipe_phi.unet,
+                torch.cat(
+                    [
+                        torch.cat([latents_noisy] * 2, dim=0),
+                        torch.cat(
+                            [self.image_latents.repeat(batch_size, 1, 1, 1)] * 2,
+                            dim=0,
+                        ),
+                    ],
+                    dim=1,
+                ),
+                torch.cat([t] * 2, dim=0),
+                encoder_hidden_states=torch.cat([image_camera_embeddings] * 2, dim=0),
+                class_labels=torch.cat(
+                    [
+                        camera_condition.view(batch_size, -1),
+                        torch.zeros_like(camera_condition.view(batch_size, -1)),
+                    ],
+                    dim=0,
+                )
+                if self.cfg.vsd_use_camera_condition
+                else None,
+                cross_attention_kwargs={"scale": 1.0},
+                velocity_to_epsilon=self.pipe_phi.scheduler.config.prediction_type
+                == "v_prediction",
+            )
+
+        noise_pred_camera, noise_pred_uncond = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + self.cfg.vsd_guidance_scale_phi * (
+            noise_pred_camera - noise_pred_uncond
+        )
+
+        return noise_pred
+
+    def train_phi(
+        self,
+        latents: Float[Tensor, "B 4 Hl Wl"],
+        image_camera_embeddings: Float[Tensor, "B 1 D"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        camera_condition: Float[Tensor, "B ..."],
+    ):
+        B = latents.shape[0]
+        latents = latents.detach().repeat(
+            self.cfg.vsd_lora_n_timestamp_samples, 1, 1, 1
+        )
+
+        num_train_timesteps = self.pipe_phi.scheduler.config.num_train_timesteps
+        t = torch.randint(
+            int(num_train_timesteps * 0.0),
+            int(num_train_timesteps * 1.0),
+            [B * self.cfg.vsd_lora_n_timestamp_samples],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        noise = torch.randn_like(latents)
+        latents_noisy = self.pipe_phi.scheduler.add_noise(latents, noise, t)
+        if self.pipe_phi.scheduler.config.prediction_type == "epsilon":
+            target = noise
+        elif self.pipe_phi.scheduler.prediction_type == "v_prediction":
+            target = self.pipe_phi.scheduler.get_velocity(latents, noise, t)
+        else:
+            raise ValueError(
+                f"Unknown prediction type {self.pipe_phi.scheduler.prediction_type}"
+            )
+
+        if (
+            self.cfg.vsd_use_camera_condition
+            and self.cfg.vsd_lora_cfg_training
+            and random.random() < 0.1
+        ):
+            camera_condition = torch.zeros_like(camera_condition)
+
+        noise_pred = self.forward_unet(
+            self.pipe_phi.unet,
+            torch.cat([latents_noisy, self.image_latents.repeat(B, 1, 1, 1)], dim=1),
+            t,
+            encoder_hidden_states=image_camera_embeddings.repeat(
+                self.cfg.vsd_lora_n_timestamp_samples, 1, 1
+            ),
+            class_labels=camera_condition.view(B, -1).repeat(
+                self.cfg.vsd_lora_n_timestamp_samples, 1
+            )
+            if self.cfg.vsd_use_camera_condition
+            else None,
+            cross_attention_kwargs={"scale": 1.0},
+        )
+        return F.mse_loss(noise_pred.float(), target.float(), reduction="mean")
+
+    def forward(
+        self,
+        rgb: Float[Tensor, "B H W C"],
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        mvp_mtx: Float[Tensor, "B 4 4"],
+        c2w: Float[Tensor, "B 4 4"],
+        rgb_as_latents=False,
+        **kwargs,
+    ):
+        batch_size = rgb.shape[0]
+
+        rgb_BCHW = rgb.permute(0, 3, 1, 2)
+        latents: Float[Tensor, "B 4 32 32"]
+        if rgb_as_latents:
+            # treat input rgb as latents
+            # input rgb should be in range [-1, 1]
+            latents = F.interpolate(
+                rgb_BCHW, (32, 32), mode="bilinear", align_corners=False
+            )
+        else:
+            # treat input rgb as rgb
+            # input rgb should be in range [0, 1]
+            rgb_BCHW = F.interpolate(
+                rgb_BCHW, (256, 256), mode="bilinear", align_corners=False
+            )
+            # encode image into latents with vae
+            latents = self.vae_encode(self.pipe.vae, rgb_BCHW * 2.0 - 1.0)
+
+        # sample timestep
+        # use the same timestep for each batch
+        assert self.min_step is not None and self.max_step is not None
+        t = torch.randint(
+            self.min_step,
+            self.max_step + 1,
+            [1],
+            dtype=torch.long,
+            device=self.device,
+        ).repeat(batch_size)
+
+        # sample noise
+        noise = torch.randn_like(latents)
+        latents_noisy = self.scheduler.add_noise(latents, noise, t)
+
+        # image-camera feature condition
+        image_camera_embeddings = self.get_image_camera_embeddings(
+            elevation, azimuth, camera_distances
+        )
+
+        eps_pretrain = self.get_eps_pretrain(
+            latents_noisy,
+            t,
+            image_camera_embeddings,
+            elevation,
+            azimuth,
+            camera_distances,
+        )
+
+        latents_1step_orig = (
+            1
+            / self.alphas[t].view(-1, 1, 1, 1)
+            * (latents_noisy - self.sigmas[t].view(-1, 1, 1, 1) * eps_pretrain)
+        ).detach()
+
+        if self.cfg.guidance_type == "sds":
+            eps_phi = noise
+        elif self.cfg.guidance_type == "vsd":
+            if self.cfg.vsd_camera_condition_type == "extrinsics":
+                camera_condition = c2w
+            elif self.cfg.vsd_camera_condition_type == "mvp":
+                camera_condition = mvp_mtx
+            elif self.cfg.vsd_camera_condition_type == "spherical":
+                camera_condition = torch.stack(
+                    [
+                        torch.deg2rad(elevation),
+                        torch.sin(torch.deg2rad(azimuth)),
+                        torch.cos(torch.deg2rad(azimuth)),
+                        camera_distances,
+                    ],
+                    dim=-1,
+                )
+            else:
+                raise ValueError(
+                    f"Unknown camera_condition_type {self.cfg.vsd_camera_condition_type}"
+                )
+            eps_phi = self.get_eps_phi(
+                latents_noisy,
+                t,
+                image_camera_embeddings,
+                elevation,
+                azimuth,
+                camera_distances,
+                camera_condition,
+            )
+
+            loss_train_phi = self.train_phi(
+                latents,
+                image_camera_embeddings,
+                elevation,
+                azimuth,
+                camera_distances,
+                camera_condition,
+            )
+
+        if self.cfg.weighting_strategy == "dreamfusion":
+            w = (1.0 - self.alphas[t]).view(-1, 1, 1, 1)
+        elif self.cfg.weighting_strategy == "uniform":
+            w = 1.0
+        elif self.cfg.weighting_strategy == "fantasia3d":
+            w = (self.alphas[t] ** 0.5 * (1 - self.alphas[t])).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(
+                f"Unknown weighting strategy: {self.cfg.weighting_strategy}"
+            )
+
+        grad = w * (eps_pretrain - eps_phi)
+
+        if self.grad_clip_val is not None:
+            grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val)
+
+        # reparameterization trick:
+        # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad
+        target = (latents - grad).detach()
+        loss_sd = 0.5 * F.mse_loss(latents, target, reduction="sum") / batch_size
+
+        guidance_out = {
+            "loss_sd": loss_sd,
+            "grad_norm": grad.norm(),
+            "timesteps": t,
+            "min_step": self.min_step,
+            "max_step": self.max_step,
+            "latents": latents,
+            "latents_1step_orig": latents_1step_orig,
+            "rgb": rgb_BCHW.permute(0, 2, 3, 1),
+            "weights": w,
+            "lambdas": self.lambdas[t],
+        }
+
+        if self.cfg.return_rgb_1step_orig:
+            with torch.no_grad():
+                rgb_1step_orig = self.vae_decode(
+                    self.pipe.vae, latents_1step_orig
+                ).permute(0, 2, 3, 1)
+            guidance_out.update({"rgb_1step_orig": rgb_1step_orig})
+
+        if self.cfg.return_rgb_multistep_orig:
+            with self.set_scheduler(
+                self.pipe,
+                DPMSolverSinglestepScheduler,
+                solver_order=1,
+                num_train_timesteps=int(t[0]),
+            ) as pipe:
+                with torch.cuda.amp.autocast(enabled=False):
+                    latents_multistep_orig = pipe(
+                        num_inference_steps=self.cfg.n_rgb_multistep_orig_steps,
+                        guidance_scale=self.cfg.guidance_scale,
+                        eta=1.0,
+                        latents=latents_noisy.to(pipe.unet.dtype),
+                        image_camera_embeddings=image_camera_embeddings.to(
+                            pipe.unet.dtype
+                        ),
+                        image_latents=self.image_latents.repeat(batch_size, 1, 1, 1).to(
+                            pipe.unet.dtype
+                        ),
+                        cross_attention_kwargs={"scale": 0.0}
+                        if self.vsd_share_model
+                        else None,
+                        output_type="latent",
+                    ).images.to(latents.dtype)
+            with torch.no_grad():
+                rgb_multistep_orig = self.vae_decode(
+                    self.pipe.vae, latents_multistep_orig
+                )
+            guidance_out.update(
+                {
+                    "latents_multistep_orig": latents_multistep_orig,
+                    "rgb_multistep_orig": rgb_multistep_orig.permute(0, 2, 3, 1),
+                }
+            )
+
+        if self.cfg.guidance_type == "vsd":
+            guidance_out.update(
+                {
+                    "loss_train_phi": loss_train_phi,
+                }
+            )
+
+        return guidance_out
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # clip grad for stable training as demonstrated in
+        # Debiasing Scores and Prompts of 2D Diffusion for Robust Text-to-3D Generation
+        # http://arxiv.org/abs/2303.15413
+        if self.cfg.grad_clip is not None:
+            self.grad_clip_val = C(self.cfg.grad_clip, epoch, global_step)
+
+        self.min_step = int(
+            self.num_train_timesteps * C(self.cfg.min_step_percent, epoch, global_step)
+        )
+        self.max_step = int(
+            self.num_train_timesteps * C(self.cfg.max_step_percent, epoch, global_step)
+        )

threestudio/models/isosurface.py

@@ -0,0 +1,253 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.mesh import Mesh
+from threestudio.utils.typing import *
+
+
+class IsosurfaceHelper(nn.Module):
+    points_range: Tuple[float, float] = (0, 1)
+
+    @property
+    def grid_vertices(self) -> Float[Tensor, "N 3"]:
+        raise NotImplementedError
+
+
+class MarchingCubeCPUHelper(IsosurfaceHelper):
+    def __init__(self, resolution: int) -> None:
+        super().__init__()
+        self.resolution = resolution
+        import mcubes
+
+        self.mc_func: Callable = mcubes.marching_cubes
+        self._grid_vertices: Optional[Float[Tensor, "N3 3"]] = None
+        self._dummy: Float[Tensor, "..."]
+        self.register_buffer(
+            "_dummy", torch.zeros(0, dtype=torch.float32), persistent=False
+        )
+
+    @property
+    def grid_vertices(self) -> Float[Tensor, "N3 3"]:
+        if self._grid_vertices is None:
+            # keep the vertices on CPU so that we can support very large resolution
+            x, y, z = (
+                torch.linspace(*self.points_range, self.resolution),
+                torch.linspace(*self.points_range, self.resolution),
+                torch.linspace(*self.points_range, self.resolution),
+            )
+            x, y, z = torch.meshgrid(x, y, z, indexing="ij")
+            verts = torch.cat(
+                [x.reshape(-1, 1), y.reshape(-1, 1), z.reshape(-1, 1)], dim=-1
+            ).reshape(-1, 3)
+            self._grid_vertices = verts
+        return self._grid_vertices
+
+    def forward(
+        self,
+        level: Float[Tensor, "N3 1"],
+        deformation: Optional[Float[Tensor, "N3 3"]] = None,
+    ) -> Mesh:
+        if deformation is not None:
+            threestudio.warn(
+                f"{self.__class__.__name__} does not support deformation. Ignoring."
+            )
+        level = -level.view(self.resolution, self.resolution, self.resolution)
+        v_pos, t_pos_idx = self.mc_func(
+            level.detach().cpu().numpy(), 0.0
+        )  # transform to numpy
+        v_pos, t_pos_idx = (
+            torch.from_numpy(v_pos).float().to(self._dummy.device),
+            torch.from_numpy(t_pos_idx.astype(np.int64)).long().to(self._dummy.device),
+        )  # transform back to torch tensor on CUDA
+        v_pos = v_pos / (self.resolution - 1.0)
+        return Mesh(v_pos=v_pos, t_pos_idx=t_pos_idx)
+
+
+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
+
+    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

threestudio/models/materials/__init__.py

@@ -0,0 +1,9 @@
+from . import (
+    base,
+    diffuse_with_point_light_material,
+    hybrid_rgb_latent_material,
+    neural_radiance_material,
+    no_material,
+    pbr_material,
+    sd_latent_adapter_material,
+)

threestudio/models/materials/base.py

@@ -0,0 +1,29 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.utils.base import BaseModule
+from threestudio.utils.typing import *
+
+
+class BaseMaterial(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        pass
+
+    cfg: Config
+    requires_normal: bool = False
+    requires_tangent: bool = False
+
+    def configure(self):
+        pass
+
+    def forward(self, *args, **kwargs) -> Float[Tensor, "*B 3"]:
+        raise NotImplementedError
+
+    def export(self, *args, **kwargs) -> Dict[str, Any]:
+        return {}

threestudio/models/materials/diffuse_with_point_light_material.py

@@ -0,0 +1,120 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.utils.ops import dot, get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("diffuse-with-point-light-material")
+class DiffuseWithPointLightMaterial(BaseMaterial):
+    @dataclass
+    class Config(BaseMaterial.Config):
+        ambient_light_color: Tuple[float, float, float] = (0.1, 0.1, 0.1)
+        diffuse_light_color: Tuple[float, float, float] = (0.9, 0.9, 0.9)
+        ambient_only_steps: int = 1000
+        diffuse_prob: float = 0.75
+        textureless_prob: float = 0.5
+        albedo_activation: str = "sigmoid"
+        soft_shading: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.requires_normal = True
+
+        self.ambient_light_color: Float[Tensor, "3"]
+        self.register_buffer(
+            "ambient_light_color",
+            torch.as_tensor(self.cfg.ambient_light_color, dtype=torch.float32),
+        )
+        self.diffuse_light_color: Float[Tensor, "3"]
+        self.register_buffer(
+            "diffuse_light_color",
+            torch.as_tensor(self.cfg.diffuse_light_color, dtype=torch.float32),
+        )
+        self.ambient_only = False
+
+    def forward(
+        self,
+        features: Float[Tensor, "B ... Nf"],
+        positions: Float[Tensor, "B ... 3"],
+        shading_normal: Float[Tensor, "B ... 3"],
+        light_positions: Float[Tensor, "B ... 3"],
+        ambient_ratio: Optional[float] = None,
+        shading: Optional[str] = None,
+        **kwargs,
+    ) -> Float[Tensor, "B ... 3"]:
+        albedo = get_activation(self.cfg.albedo_activation)(features[..., :3])
+
+        if ambient_ratio is not None:
+            # if ambient ratio is specified, use it
+            diffuse_light_color = (1 - ambient_ratio) * torch.ones_like(
+                self.diffuse_light_color
+            )
+            ambient_light_color = ambient_ratio * torch.ones_like(
+                self.ambient_light_color
+            )
+        elif self.training and self.cfg.soft_shading:
+            # otherwise if in training and soft shading is enabled, random a ambient ratio
+            diffuse_light_color = torch.full_like(
+                self.diffuse_light_color, random.random()
+            )
+            ambient_light_color = 1.0 - diffuse_light_color
+        else:
+            # otherwise use the default fixed values
+            diffuse_light_color = self.diffuse_light_color
+            ambient_light_color = self.ambient_light_color
+
+        light_directions: Float[Tensor, "B ... 3"] = F.normalize(
+            light_positions - positions, dim=-1
+        )
+        diffuse_light: Float[Tensor, "B ... 3"] = (
+            dot(shading_normal, light_directions).clamp(min=0.0) * diffuse_light_color
+        )
+        textureless_color = diffuse_light + ambient_light_color
+        # clamp albedo to [0, 1] to compute shading
+        color = albedo.clamp(0.0, 1.0) * textureless_color
+
+        if shading is None:
+            if self.training:
+                # adopt the same type of augmentation for the whole batch
+                if self.ambient_only or random.random() > self.cfg.diffuse_prob:
+                    shading = "albedo"
+                elif random.random() < self.cfg.textureless_prob:
+                    shading = "textureless"
+                else:
+                    shading = "diffuse"
+            else:
+                if self.ambient_only:
+                    shading = "albedo"
+                else:
+                    # return shaded color by default in evaluation
+                    shading = "diffuse"
+
+        # multiply by 0 to prevent checking for unused parameters in DDP
+        if shading == "albedo":
+            return albedo + textureless_color * 0
+        elif shading == "textureless":
+            return albedo * 0 + textureless_color
+        elif shading == "diffuse":
+            return color
+        else:
+            raise ValueError(f"Unknown shading type {shading}")
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        if global_step < self.cfg.ambient_only_steps:
+            self.ambient_only = True
+        else:
+            self.ambient_only = False
+
+    def export(self, features: Float[Tensor, "*N Nf"], **kwargs) -> Dict[str, Any]:
+        albedo = get_activation(self.cfg.albedo_activation)(features[..., :3]).clamp(
+            0.0, 1.0
+        )
+        return {"albedo": albedo}

threestudio/models/materials/hybrid_rgb_latent_material.py

@@ -0,0 +1,36 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.ops import dot, get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("hybrid-rgb-latent-material")
+class HybridRGBLatentMaterial(BaseMaterial):
+    @dataclass
+    class Config(BaseMaterial.Config):
+        n_output_dims: int = 3
+        color_activation: str = "sigmoid"
+        requires_normal: bool = True
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.requires_normal = self.cfg.requires_normal
+
+    def forward(
+        self, features: Float[Tensor, "B ... Nf"], **kwargs
+    ) -> Float[Tensor, "B ... Nc"]:
+        assert (
+            features.shape[-1] == self.cfg.n_output_dims
+        ), f"Expected {self.cfg.n_output_dims} output dims, only got {features.shape[-1]} dims input."
+        color = features
+        color[..., :3] = get_activation(self.cfg.color_activation)(color[..., :3])
+        return color

threestudio/models/materials/neural_radiance_material.py

@@ -0,0 +1,54 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.ops import dot, get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("neural-radiance-material")
+class NeuralRadianceMaterial(BaseMaterial):
+    @dataclass
+    class Config(BaseMaterial.Config):
+        input_feature_dims: int = 8
+        color_activation: str = "sigmoid"
+        dir_encoding_config: dict = field(
+            default_factory=lambda: {"otype": "SphericalHarmonics", "degree": 3}
+        )
+        mlp_network_config: dict = field(
+            default_factory=lambda: {
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "n_neurons": 16,
+                "n_hidden_layers": 2,
+            }
+        )
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.encoding = get_encoding(3, self.cfg.dir_encoding_config)
+        self.n_input_dims = self.cfg.input_feature_dims + self.encoding.n_output_dims  # type: ignore
+        self.network = get_mlp(self.n_input_dims, 3, self.cfg.mlp_network_config)
+
+    def forward(
+        self,
+        features: Float[Tensor, "*B Nf"],
+        viewdirs: Float[Tensor, "*B 3"],
+        **kwargs,
+    ) -> Float[Tensor, "*B 3"]:
+        # viewdirs and normals must be normalized before passing to this function
+        viewdirs = (viewdirs + 1.0) / 2.0  # (-1, 1) => (0, 1)
+        viewdirs_embd = self.encoding(viewdirs.view(-1, 3))
+        network_inp = torch.cat(
+            [features.view(-1, features.shape[-1]), viewdirs_embd], dim=-1
+        )
+        color = self.network(network_inp).view(*features.shape[:-1], 3)
+        color = get_activation(self.cfg.color_activation)(color)
+        return color

threestudio/models/materials/no_material.py

@@ -0,0 +1,63 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.ops import dot, get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("no-material")
+class NoMaterial(BaseMaterial):
+    @dataclass
+    class Config(BaseMaterial.Config):
+        n_output_dims: int = 3
+        color_activation: str = "sigmoid"
+        input_feature_dims: Optional[int] = None
+        mlp_network_config: Optional[dict] = None
+        requires_normal: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.use_network = False
+        if (
+            self.cfg.input_feature_dims is not None
+            and self.cfg.mlp_network_config is not None
+        ):
+            self.network = get_mlp(
+                self.cfg.input_feature_dims,
+                self.cfg.n_output_dims,
+                self.cfg.mlp_network_config,
+            )
+            self.use_network = True
+        self.requires_normal = self.cfg.requires_normal
+
+    def forward(
+        self, features: Float[Tensor, "B ... Nf"], **kwargs
+    ) -> Float[Tensor, "B ... Nc"]:
+        if not self.use_network:
+            assert (
+                features.shape[-1] == self.cfg.n_output_dims
+            ), f"Expected {self.cfg.n_output_dims} output dims, only got {features.shape[-1]} dims input."
+            color = get_activation(self.cfg.color_activation)(features)
+        else:
+            color = self.network(features.view(-1, features.shape[-1])).view(
+                *features.shape[:-1], self.cfg.n_output_dims
+            )
+            color = get_activation(self.cfg.color_activation)(color)
+        return color
+
+    def export(self, features: Float[Tensor, "*N Nf"], **kwargs) -> Dict[str, Any]:
+        color = self(features, **kwargs).clamp(0, 1)
+        assert color.shape[-1] >= 3, "Output color must have at least 3 channels"
+        if color.shape[-1] > 3:
+            threestudio.warn(
+                "Output color has >3 channels, treating the first 3 as RGB"
+            )
+        return {"albedo": color[..., :3]}

threestudio/models/materials/pbr_material.py

@@ -0,0 +1,143 @@
+import random
+from dataclasses import dataclass, field
+
+import envlight
+import numpy as np
+import nvdiffrast.torch as dr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.utils.ops import get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("pbr-material")
+class PBRMaterial(BaseMaterial):
+    @dataclass
+    class Config(BaseMaterial.Config):
+        material_activation: str = "sigmoid"
+        environment_texture: str = "load/lights/mud_road_puresky_1k.hdr"
+        environment_scale: float = 2.0
+        min_metallic: float = 0.0
+        max_metallic: float = 0.9
+        min_roughness: float = 0.08
+        max_roughness: float = 0.9
+        use_bump: bool = True
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.requires_normal = True
+        self.requires_tangent = self.cfg.use_bump
+
+        self.light = envlight.EnvLight(
+            self.cfg.environment_texture, scale=self.cfg.environment_scale
+        )
+
+        FG_LUT = torch.from_numpy(
+            np.fromfile("load/lights/bsdf_256_256.bin", dtype=np.float32).reshape(
+                1, 256, 256, 2
+            )
+        )
+        self.register_buffer("FG_LUT", FG_LUT)
+
+    def forward(
+        self,
+        features: Float[Tensor, "*B Nf"],
+        viewdirs: Float[Tensor, "*B 3"],
+        shading_normal: Float[Tensor, "B ... 3"],
+        tangent: Optional[Float[Tensor, "B ... 3"]] = None,
+        **kwargs,
+    ) -> Float[Tensor, "*B 3"]:
+        prefix_shape = features.shape[:-1]
+
+        material: Float[Tensor, "*B Nf"] = get_activation(self.cfg.material_activation)(
+            features
+        )
+        albedo = material[..., :3]
+        metallic = (
+            material[..., 3:4] * (self.cfg.max_metallic - self.cfg.min_metallic)
+            + self.cfg.min_metallic
+        )
+        roughness = (
+            material[..., 4:5] * (self.cfg.max_roughness - self.cfg.min_roughness)
+            + self.cfg.min_roughness
+        )
+
+        if self.cfg.use_bump:
+            assert tangent is not None
+            # perturb_normal is a delta to the initialization [0, 0, 1]
+            perturb_normal = (material[..., 5:8] * 2 - 1) + torch.tensor(
+                [0, 0, 1], dtype=material.dtype, device=material.device
+            )
+            perturb_normal = F.normalize(perturb_normal.clamp(-1, 1), dim=-1)
+
+            # apply normal perturbation in tangent space
+            bitangent = F.normalize(torch.cross(tangent, shading_normal), dim=-1)
+            shading_normal = (
+                tangent * perturb_normal[..., 0:1]
+                - bitangent * perturb_normal[..., 1:2]
+                + shading_normal * perturb_normal[..., 2:3]
+            )
+            shading_normal = F.normalize(shading_normal, dim=-1)
+
+        v = -viewdirs
+        n_dot_v = (shading_normal * v).sum(-1, keepdim=True)
+        reflective = n_dot_v * shading_normal * 2 - v
+
+        diffuse_albedo = (1 - metallic) * albedo
+
+        fg_uv = torch.cat([n_dot_v, roughness], -1).clamp(0, 1)
+        fg = dr.texture(
+            self.FG_LUT,
+            fg_uv.reshape(1, -1, 1, 2).contiguous(),
+            filter_mode="linear",
+            boundary_mode="clamp",
+        ).reshape(*prefix_shape, 2)
+        F0 = (1 - metallic) * 0.04 + metallic * albedo
+        specular_albedo = F0 * fg[:, 0:1] + fg[:, 1:2]
+
+        diffuse_light = self.light(shading_normal)
+        specular_light = self.light(reflective, roughness)
+
+        color = diffuse_albedo * diffuse_light + specular_albedo * specular_light
+        color = color.clamp(0.0, 1.0)
+
+        return color
+
+    def export(self, features: Float[Tensor, "*N Nf"], **kwargs) -> Dict[str, Any]:
+        material: Float[Tensor, "*N Nf"] = get_activation(self.cfg.material_activation)(
+            features
+        )
+        albedo = material[..., :3]
+        metallic = (
+            material[..., 3:4] * (self.cfg.max_metallic - self.cfg.min_metallic)
+            + self.cfg.min_metallic
+        )
+        roughness = (
+            material[..., 4:5] * (self.cfg.max_roughness - self.cfg.min_roughness)
+            + self.cfg.min_roughness
+        )
+
+        out = {
+            "albedo": albedo,
+            "metallic": metallic,
+            "roughness": roughness,
+        }
+
+        if self.cfg.use_bump:
+            perturb_normal = (material[..., 5:8] * 2 - 1) + torch.tensor(
+                [0, 0, 1], dtype=material.dtype, device=material.device
+            )
+            perturb_normal = F.normalize(perturb_normal.clamp(-1, 1), dim=-1)
+            perturb_normal = (perturb_normal + 1) / 2
+            out.update(
+                {
+                    "bump": perturb_normal,
+                }
+            )
+
+        return out

threestudio/models/materials/sd_latent_adapter_material.py

@@ -0,0 +1,42 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.utils.typing import *
+
+
+@threestudio.register("sd-latent-adapter-material")
+class StableDiffusionLatentAdapterMaterial(BaseMaterial):
+    @dataclass
+    class Config(BaseMaterial.Config):
+        pass
+
+    cfg: Config
+
+    def configure(self) -> None:
+        adapter = nn.Parameter(
+            torch.as_tensor(
+                [
+                    #   R       G       B
+                    [0.298, 0.207, 0.208],  # L1
+                    [0.187, 0.286, 0.173],  # L2
+                    [-0.158, 0.189, 0.264],  # L3
+                    [-0.184, -0.271, -0.473],  # L4
+                ]
+            )
+        )
+        self.register_parameter("adapter", adapter)
+
+    def forward(
+        self, features: Float[Tensor, "B ... 4"], **kwargs
+    ) -> Float[Tensor, "B ... 3"]:
+        assert features.shape[-1] == 4
+        color = features @ self.adapter
+        color = (color + 1) / 2
+        color = color.clamp(0.0, 1.0)
+        return color

threestudio/models/mesh.py

@@ -0,0 +1,309 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.utils.ops import dot
+from threestudio.utils.typing import *
+
+
+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._t_tex_idx: Optional[Float[Tensor, "Nf 3"]] = None
+        self._v_rgb: Optional[Float[Tensor, "Nv 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)
+
+    def add_extra(self, k, v) -> None:
+        self.extras[k] = v
+
+    def remove_outlier(self, outlier_n_faces_threshold: Union[int, float]) -> Mesh:
+        if self.requires_grad:
+            threestudio.debug("Mesh is differentiable, not removing outliers")
+            return self
+
+        # use trimesh to first split the mesh into connected components
+        # then remove the components with less than n_face_threshold faces
+        import trimesh
+
+        # construct a trimesh object
+        mesh = trimesh.Trimesh(
+            vertices=self.v_pos.detach().cpu().numpy(),
+            faces=self.t_pos_idx.detach().cpu().numpy(),
+        )
+
+        # split the mesh into connected components
+        components = mesh.split(only_watertight=False)
+        # log the number of faces in each component
+        threestudio.debug(
+            "Mesh has {} components, with faces: {}".format(
+                len(components), [c.faces.shape[0] for c in components]
+            )
+        )
+
+        n_faces_threshold: int
+        if isinstance(outlier_n_faces_threshold, float):
+            # set the threshold to the number of faces in the largest component multiplied by outlier_n_faces_threshold
+            n_faces_threshold = int(
+                max([c.faces.shape[0] for c in components]) * outlier_n_faces_threshold
+            )
+        else:
+            # set the threshold directly to outlier_n_faces_threshold
+            n_faces_threshold = outlier_n_faces_threshold
+
+        # log the threshold
+        threestudio.debug(
+            "Removing components with less than {} faces".format(n_faces_threshold)
+        )
+
+        # remove the components with less than n_face_threshold faces
+        components = [c for c in components if c.faces.shape[0] >= n_faces_threshold]
+
+        # log the number of faces in each component after removing outliers
+        threestudio.debug(
+            "Mesh has {} components after removing outliers, with faces: {}".format(
+                len(components), [c.faces.shape[0] for c in components]
+            )
+        )
+        # merge the components
+        mesh = trimesh.util.concatenate(components)
+
+        # convert back to our mesh format
+        v_pos = torch.from_numpy(mesh.vertices).to(self.v_pos)
+        t_pos_idx = torch.from_numpy(mesh.faces).to(self.t_pos_idx)
+
+        clean_mesh = Mesh(v_pos, t_pos_idx)
+        # keep the extras unchanged
+
+        if len(self.extras) > 0:
+            clean_mesh.extras = self.extras
+            threestudio.debug(
+                f"The following extra attributes are inherited from the original mesh unchanged: {list(self.extras.keys())}"
+            )
+        return clean_mesh
+
+    @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._v_tex, self._t_tex_idx = self._unwrap_uv()
+        return self._v_tex
+
+    @property
+    def t_tex_idx(self):
+        if self._t_tex_idx is None:
+            self._v_tex, self._t_tex_idx = self._unwrap_uv()
+        return self._t_tex_idx
+
+    @property
+    def v_rgb(self):
+        return self._v_rgb
+
+    @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)
+
+        # 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_tex_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
+        uve1 = tex[1] - tex[0]
+        uve2 = tex[2] - tex[0]
+        pe1 = pos[1] - pos[0]
+        pe2 = pos[2] - pos[0]
+
+        nom = pe1 * uve2[..., 1:2] - pe2 * uve1[..., 1:2]
+        denom = uve1[..., 0:1] * uve2[..., 1:2] - uve1[..., 1:2] * uve2[..., 0:1]
+
+        # Avoid division by zero for degenerated texture coordinates
+        tang = nom / torch.where(
+            denom > 0.0, torch.clamp(denom, min=1e-6), torch.clamp(denom, max=-1e-6)
+        )
+
+        # Update all 3 vertices
+        for i in range(0, 3):
+            idx = vn_idx[i][:, None].repeat(1, 3)
+            tangents.scatter_add_(0, idx, tang)  # tangents[n_i] = tangents[n_i] + tang
+            tansum.scatter_add_(
+                0, idx, torch.ones_like(tang)
+            )  # tansum[n_i] = tansum[n_i] + 1
+        tangents = tangents / tansum
+
+        # Normalize and make sure tangent is perpendicular to normal
+        tangents = 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 _unwrap_uv(
+        self, xatlas_chart_options: dict = {}, xatlas_pack_options: dict = {}
+    ):
+        threestudio.info("Using xatlas to perform UV unwrapping, may take a while ...")
+
+        import xatlas
+
+        atlas = xatlas.Atlas()
+        atlas.add_mesh(
+            self.v_pos.detach().cpu().numpy(),
+            self.t_pos_idx.cpu().numpy(),
+        )
+        co = xatlas.ChartOptions()
+        po = xatlas.PackOptions()
+        for k, v in xatlas_chart_options.items():
+            setattr(co, k, v)
+        for k, v in xatlas_pack_options.items():
+            setattr(po, k, v)
+        atlas.generate(co, po)
+        vmapping, indices, uvs = atlas.get_mesh(0)
+        vmapping = (
+            torch.from_numpy(
+                vmapping.astype(np.uint64, casting="same_kind").view(np.int64)
+            )
+            .to(self.v_pos.device)
+            .long()
+        )
+        uvs = torch.from_numpy(uvs).to(self.v_pos.device).float()
+        indices = (
+            torch.from_numpy(
+                indices.astype(np.uint64, casting="same_kind").view(np.int64)
+            )
+            .to(self.v_pos.device)
+            .long()
+        )
+        return uvs, indices
+
+    def unwrap_uv(
+        self, xatlas_chart_options: dict = {}, xatlas_pack_options: dict = {}
+    ):
+        self._v_tex, self._t_tex_idx = self._unwrap_uv(
+            xatlas_chart_options, xatlas_pack_options
+        )
+
+    def set_vertex_color(self, v_rgb):
+        assert v_rgb.shape[0] == self.v_pos.shape[0]
+        self._v_rgb = v_rgb
+
+    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
+
+    def normal_consistency(self) -> Float[Tensor, ""]:
+        edge_nrm: Float[Tensor, "Ne 2 3"] = self.v_nrm[self.edges]
+        nc = (
+            1.0 - torch.cosine_similarity(edge_nrm[:, 0], edge_nrm[:, 1], dim=-1)
+        ).mean()
+        return nc
+
+    def _laplacian_uniform(self):
+        # from stable-dreamfusion
+        # https://github.com/ashawkey/stable-dreamfusion/blob/8fb3613e9e4cd1ded1066b46e80ca801dfb9fd06/nerf/renderer.py#L224
+        verts, faces = self.v_pos, self.t_pos_idx
+
+        V = verts.shape[0]
+        F = faces.shape[0]
+
+        # Neighbor indices
+        ii = faces[:, [1, 2, 0]].flatten()
+        jj = faces[:, [2, 0, 1]].flatten()
+        adj = torch.stack([torch.cat([ii, jj]), torch.cat([jj, ii])], dim=0).unique(
+            dim=1
+        )
+        adj_values = torch.ones(adj.shape[1]).to(verts)
+
+        # Diagonal indices
+        diag_idx = adj[0]
+
+        # Build the sparse matrix
+        idx = torch.cat((adj, torch.stack((diag_idx, diag_idx), dim=0)), dim=1)
+        values = torch.cat((-adj_values, adj_values))
+
+        # The coalesce operation sums the duplicate indices, resulting in the
+        # correct diagonal
+        return torch.sparse_coo_tensor(idx, values, (V, V)).coalesce()
+
+    def laplacian(self) -> Float[Tensor, ""]:
+        with torch.no_grad():
+            L = self._laplacian_uniform()
+        loss = L.mm(self.v_pos)
+        loss = loss.norm(dim=1)
+        loss = loss.mean()
+        return loss

threestudio/models/networks.py

@@ -0,0 +1,411 @@
+import math
+
+import tinycudann as tcnn
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.utils.base import Updateable
+from threestudio.utils.config import config_to_primitive
+from threestudio.utils.misc import get_rank
+from threestudio.utils.ops import get_activation
+from threestudio.utils.typing import *
+
+
+class ProgressiveBandFrequency(nn.Module, Updateable):
+    def __init__(self, in_channels: int, config: dict):
+        super().__init__()
+        self.N_freqs = config["n_frequencies"]
+        self.in_channels, self.n_input_dims = in_channels, in_channels
+        self.funcs = [torch.sin, torch.cos]
+        self.freq_bands = 2 ** torch.linspace(0, self.N_freqs - 1, self.N_freqs)
+        self.n_output_dims = self.in_channels * (len(self.funcs) * self.N_freqs)
+        self.n_masking_step = config.get("n_masking_step", 0)
+        self.update_step(
+            None, None
+        )  # mask should be updated at the beginning each step
+
+    def forward(self, x):
+        out = []
+        for freq, mask in zip(self.freq_bands, self.mask):
+            for func in self.funcs:
+                out += [func(freq * x) * mask]
+        return torch.cat(out, -1)
+
+    def update_step(self, epoch, global_step, on_load_weights=False):
+        if self.n_masking_step <= 0 or global_step is None:
+            self.mask = torch.ones(self.N_freqs, dtype=torch.float32)
+        else:
+            self.mask = (
+                1.0
+                - torch.cos(
+                    math.pi
+                    * (
+                        global_step / self.n_masking_step * self.N_freqs
+                        - torch.arange(0, self.N_freqs)
+                    ).clamp(0, 1)
+                )
+            ) / 2.0
+            threestudio.debug(
+                f"Update mask: {global_step}/{self.n_masking_step} {self.mask}"
+            )
+
+
+class TCNNEncoding(nn.Module):
+    def __init__(self, in_channels, config, dtype=torch.float32) -> None:
+        super().__init__()
+        self.n_input_dims = in_channels
+        with torch.cuda.device(get_rank()):
+            self.encoding = tcnn.Encoding(in_channels, config, dtype=dtype)
+        self.n_output_dims = self.encoding.n_output_dims
+
+    def forward(self, x):
+        return self.encoding(x)
+
+
+# 4D implicit decomposition of space and time (4D-fy)
+class TCNNEncodingSpatialTime(nn.Module):
+    def __init__(
+        self, in_channels, config, dtype=torch.float32, init_time_zero=False
+    ) -> None:
+        super().__init__()
+        self.n_input_dims = in_channels
+        config["otype"] = "HashGrid"
+        self.num_frames = 1  # config["num_frames"]
+        self.static = config["static"]
+        self.cfg = config_to_primitive(config)
+        self.cfg_time = self.cfg
+        self.n_key_frames = config.get("n_key_frames", 1)
+        with torch.cuda.device(get_rank()):
+            self.encoding = tcnn.Encoding(self.n_input_dims, self.cfg, dtype=dtype)
+            self.encoding_time = tcnn.Encoding(
+                self.n_input_dims + 1, self.cfg_time, dtype=dtype
+            )
+        self.n_output_dims = self.encoding.n_output_dims
+        self.frame_time = None
+        if self.static:
+            self.set_temp_param_grad(requires_grad=False)
+        self.use_key_frame = config.get("use_key_frame", False)
+        self.is_video = True
+        self.update_occ_grid = False
+
+    def set_temp_param_grad(self, requires_grad=False):
+        self.set_param_grad(self.encoding_time, requires_grad=requires_grad)
+
+    def set_param_grad(self, param_list, requires_grad=False):
+        if isinstance(param_list, nn.Parameter):
+            param_list.requires_grad = requires_grad
+        else:
+            for param in param_list.parameters():
+                param.requires_grad = requires_grad
+
+    def forward(self, x):
+        # TODO frame_time only supports batch_size == 1 cases
+        if self.update_occ_grid and not isinstance(self.frame_time, float):
+            frame_time = self.frame_time
+        else:
+            if (self.static or not self.training) and self.frame_time is None:
+                frame_time = torch.zeros(
+                    (self.num_frames, 1), device=x.device, dtype=x.dtype
+                ).expand(x.shape[0], 1)
+            else:
+                if self.frame_time is None:
+                    frame_time = 0.0
+                else:
+                    frame_time = self.frame_time
+                frame_time = (
+                    torch.ones((self.num_frames, 1), device=x.device, dtype=x.dtype)
+                    * frame_time
+                ).expand(x.shape[0], 1)
+            frame_time = frame_time.view(-1, 1)
+        enc_space = self.encoding(x)
+        x_frame_time = torch.cat((x, frame_time), 1)
+        enc_space_time = self.encoding_time(x_frame_time)
+        enc = enc_space + enc_space_time
+        return enc
+
+
+class ProgressiveBandHashGrid(nn.Module, Updateable):
+    def __init__(self, in_channels, config, dtype=torch.float32):
+        super().__init__()
+        self.n_input_dims = in_channels
+        encoding_config = config.copy()
+        encoding_config["otype"] = "Grid"
+        encoding_config["type"] = "Hash"
+        with torch.cuda.device(get_rank()):
+            self.encoding = tcnn.Encoding(in_channels, encoding_config, dtype=dtype)
+        self.n_output_dims = self.encoding.n_output_dims
+        self.n_level = config["n_levels"]
+        self.n_features_per_level = config["n_features_per_level"]
+        self.start_level, self.start_step, self.update_steps = (
+            config["start_level"],
+            config["start_step"],
+            config["update_steps"],
+        )
+        self.current_level = self.start_level
+        self.mask = torch.zeros(
+            self.n_level * self.n_features_per_level,
+            dtype=torch.float32,
+            device=get_rank(),
+        )
+
+    def forward(self, x):
+        enc = self.encoding(x)
+        enc = enc * self.mask
+        return enc
+
+    def update_step(self, epoch, global_step, on_load_weights=False):
+        current_level = min(
+            self.start_level
+            + max(global_step - self.start_step, 0) // self.update_steps,
+            self.n_level,
+        )
+        if current_level > self.current_level:
+            threestudio.debug(f"Update current level to {current_level}")
+        self.current_level = current_level
+        self.mask[: self.current_level * self.n_features_per_level] = 1.0
+
+
+class CompositeEncoding(nn.Module, Updateable):
+    def __init__(self, encoding, include_xyz=False, xyz_scale=2.0, xyz_offset=-1.0):
+        super(CompositeEncoding, self).__init__()
+        self.encoding = encoding
+        self.include_xyz, self.xyz_scale, self.xyz_offset = (
+            include_xyz,
+            xyz_scale,
+            xyz_offset,
+        )
+        self.n_output_dims = (
+            int(self.include_xyz) * self.encoding.n_input_dims
+            + self.encoding.n_output_dims
+        )
+
+    def forward(self, x, *args):
+        return (
+            self.encoding(x, *args)
+            if not self.include_xyz
+            else torch.cat(
+                [x * self.xyz_scale + self.xyz_offset, self.encoding(x, *args)], dim=-1
+            )
+        )
+
+
+def get_encoding(n_input_dims: int, config) -> nn.Module:
+    # input suppose to be range [0, 1]
+    encoding: nn.Module
+    if config.otype == "ProgressiveBandFrequency":
+        encoding = ProgressiveBandFrequency(n_input_dims, config_to_primitive(config))
+    elif config.otype == "ProgressiveBandHashGrid":
+        encoding = ProgressiveBandHashGrid(n_input_dims, config_to_primitive(config))
+    elif config.otype == "HashGridSpatialTime":
+        encoding = TCNNEncodingSpatialTime(n_input_dims, config)  # 4D-fy encoding
+    else:
+        encoding = TCNNEncoding(n_input_dims, config_to_primitive(config))
+    encoding = CompositeEncoding(
+        encoding,
+        include_xyz=config.get("include_xyz", False),
+        xyz_scale=2.0,
+        xyz_offset=-1.0,
+    )  # FIXME: hard coded
+    return encoding
+
+
+class VanillaMLP(nn.Module):
+    def __init__(self, dim_in: int, dim_out: int, config: dict):
+        super().__init__()
+        self.n_neurons, self.n_hidden_layers = (
+            config["n_neurons"],
+            config["n_hidden_layers"],
+        )
+        layers = [
+            self.make_linear(dim_in, self.n_neurons, is_first=True, is_last=False),
+            self.make_activation(),
+        ]
+        for i in range(self.n_hidden_layers - 1):
+            layers += [
+                self.make_linear(
+                    self.n_neurons, self.n_neurons, is_first=False, is_last=False
+                ),
+                self.make_activation(),
+            ]
+        layers += [
+            self.make_linear(self.n_neurons, dim_out, is_first=False, is_last=True)
+        ]
+        self.layers = nn.Sequential(*layers)
+        self.output_activation = get_activation(config.get("output_activation", None))
+
+    def forward(self, x):
+        # disable autocast
+        # strange that the parameters will have empty gradients if autocast is enabled in AMP
+        with torch.cuda.amp.autocast(enabled=False):
+            x = self.layers(x)
+            x = self.output_activation(x)
+        return x
+
+    def make_linear(self, dim_in, dim_out, is_first, is_last):
+        layer = nn.Linear(dim_in, dim_out, bias=False)
+        return layer
+
+    def make_activation(self):
+        return nn.ReLU(inplace=True)
+
+
+class SphereInitVanillaMLP(nn.Module):
+    def __init__(self, dim_in, dim_out, config):
+        super().__init__()
+        self.n_neurons, self.n_hidden_layers = (
+            config["n_neurons"],
+            config["n_hidden_layers"],
+        )
+        self.sphere_init, self.weight_norm = True, True
+        self.sphere_init_radius = config["sphere_init_radius"]
+        self.sphere_init_inside_out = config["inside_out"]
+
+        self.layers = [
+            self.make_linear(dim_in, self.n_neurons, is_first=True, is_last=False),
+            self.make_activation(),
+        ]
+        for i in range(self.n_hidden_layers - 1):
+            self.layers += [
+                self.make_linear(
+                    self.n_neurons, self.n_neurons, is_first=False, is_last=False
+                ),
+                self.make_activation(),
+            ]
+        self.layers += [
+            self.make_linear(self.n_neurons, dim_out, is_first=False, is_last=True)
+        ]
+        self.layers = nn.Sequential(*self.layers)
+        self.output_activation = get_activation(config.get("output_activation", None))
+
+    def forward(self, x):
+        # disable autocast
+        # strange that the parameters will have empty gradients if autocast is enabled in AMP
+        with torch.cuda.amp.autocast(enabled=False):
+            x = self.layers(x)
+            x = self.output_activation(x)
+        return x
+
+    def make_linear(self, dim_in, dim_out, is_first, is_last):
+        layer = nn.Linear(dim_in, dim_out, bias=True)
+
+        if is_last:
+            if not self.sphere_init_inside_out:
+                torch.nn.init.constant_(layer.bias, -self.sphere_init_radius)
+                torch.nn.init.normal_(
+                    layer.weight,
+                    mean=math.sqrt(math.pi) / math.sqrt(dim_in),
+                    std=0.0001,
+                )
+            else:
+                torch.nn.init.constant_(layer.bias, self.sphere_init_radius)
+                torch.nn.init.normal_(
+                    layer.weight,
+                    mean=-math.sqrt(math.pi) / math.sqrt(dim_in),
+                    std=0.0001,
+                )
+        elif is_first:
+            torch.nn.init.constant_(layer.bias, 0.0)
+            torch.nn.init.constant_(layer.weight[:, 3:], 0.0)
+            torch.nn.init.normal_(
+                layer.weight[:, :3], 0.0, math.sqrt(2) / math.sqrt(dim_out)
+            )
+        else:
+            torch.nn.init.constant_(layer.bias, 0.0)
+            torch.nn.init.normal_(layer.weight, 0.0, math.sqrt(2) / math.sqrt(dim_out))
+
+        if self.weight_norm:
+            layer = nn.utils.weight_norm(layer)
+        return layer
+
+    def make_activation(self):
+        return nn.Softplus(beta=100)
+
+
+class TCNNNetwork(nn.Module):
+    def __init__(self, dim_in: int, dim_out: int, config: dict) -> None:
+        super().__init__()
+        with torch.cuda.device(get_rank()):
+            self.network = tcnn.Network(dim_in, dim_out, config)
+
+    def forward(self, x):
+        return self.network(x).float()  # transform to float32
+
+
+def get_mlp(n_input_dims, n_output_dims, config) -> nn.Module:
+    network: nn.Module
+    if config.otype == "VanillaMLP":
+        network = VanillaMLP(n_input_dims, n_output_dims, config_to_primitive(config))
+    elif config.otype == "SphereInitVanillaMLP":
+        network = SphereInitVanillaMLP(
+            n_input_dims, n_output_dims, config_to_primitive(config)
+        )
+    else:
+        assert (
+            config.get("sphere_init", False) is False
+        ), "sphere_init=True only supported by VanillaMLP"
+        network = TCNNNetwork(n_input_dims, n_output_dims, config_to_primitive(config))
+    return network
+
+
+class NetworkWithInputEncoding(nn.Module, Updateable):
+    def __init__(self, encoding, network):
+        super().__init__()
+        self.encoding, self.network = encoding, network
+
+    def forward(self, x):
+        return self.network(self.encoding(x))
+
+
+class TCNNNetworkWithInputEncoding(nn.Module):
+    def __init__(
+        self,
+        n_input_dims: int,
+        n_output_dims: int,
+        encoding_config: dict,
+        network_config: dict,
+    ) -> None:
+        super().__init__()
+        with torch.cuda.device(get_rank()):
+            self.network_with_input_encoding = tcnn.NetworkWithInputEncoding(
+                n_input_dims=n_input_dims,
+                n_output_dims=n_output_dims,
+                encoding_config=encoding_config,
+                network_config=network_config,
+            )
+
+    def forward(self, x):
+        return self.network_with_input_encoding(x).float()  # transform to float32
+
+
+def create_network_with_input_encoding(
+    n_input_dims: int, n_output_dims: int, encoding_config, network_config
+) -> nn.Module:
+    # input suppose to be range [0, 1]
+    network_with_input_encoding: nn.Module
+    if encoding_config.otype in [
+        "VanillaFrequency",
+        "ProgressiveBandHashGrid",
+    ] or network_config.otype in ["VanillaMLP", "SphereInitVanillaMLP"]:
+        encoding = get_encoding(n_input_dims, encoding_config)
+        network = get_mlp(encoding.n_output_dims, n_output_dims, network_config)
+        network_with_input_encoding = NetworkWithInputEncoding(encoding, network)
+    else:
+        network_with_input_encoding = TCNNNetworkWithInputEncoding(
+            n_input_dims=n_input_dims,
+            n_output_dims=n_output_dims,
+            encoding_config=config_to_primitive(encoding_config),
+            network_config=config_to_primitive(network_config),
+        )
+    return network_with_input_encoding
+
+
+class ToDTypeWrapper(nn.Module):
+    def __init__(self, module: nn.Module, dtype: torch.dtype):
+        super().__init__()
+        self.module = module
+        self.dtype = dtype
+
+    def forward(self, x: Float[Tensor, "..."]) -> Float[Tensor, "..."]:
+        return self.module(x).to(self.dtype)

threestudio/models/prompt_processors/__init__.py

@@ -0,0 +1,7 @@
+from . import (
+    base,
+    deepfloyd_prompt_processor,
+    dummy_prompt_processor,
+    stable_diffusion_prompt_processor,
+    clip_prompt_processor,
+)

threestudio/models/prompt_processors/base.py

@@ -0,0 +1,517 @@
+import json
+import os
+from dataclasses import dataclass, field
+
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+import torch.nn.functional as F
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
+from transformers import AutoTokenizer, BertForMaskedLM
+
+import threestudio
+from threestudio.utils.base import BaseObject
+from threestudio.utils.misc import barrier, cleanup, get_rank
+from threestudio.utils.ops import shifted_cosine_decay, shifted_expotional_decay
+from threestudio.utils.typing import *
+
+
+def hash_prompt(model: str, prompt: str) -> str:
+    import hashlib
+
+    identifier = f"{model}-{prompt}"
+    return hashlib.md5(identifier.encode()).hexdigest()
+
+
+@dataclass
+class DirectionConfig:
+    name: str
+    prompt: Callable[[str], str]
+    negative_prompt: Callable[[str], str]
+    condition: Callable[
+        [Float[Tensor, "B"], Float[Tensor, "B"], Float[Tensor, "B"]],
+        Float[Tensor, "B"],
+    ]
+
+
+@dataclass
+class PromptProcessorOutput:
+    text_embeddings: Float[Tensor, "N Nf"]
+    uncond_text_embeddings: Float[Tensor, "N Nf"]
+    text_embeddings_vd: Float[Tensor, "Nv N Nf"]
+    uncond_text_embeddings_vd: Float[Tensor, "Nv N Nf"]
+    directions: List[DirectionConfig]
+    direction2idx: Dict[str, int]
+    use_perp_neg: bool
+    perp_neg_f_sb: Tuple[float, float, float]
+    perp_neg_f_fsb: Tuple[float, float, float]
+    perp_neg_f_fs: Tuple[float, float, float]
+    perp_neg_f_sf: Tuple[float, float, float]
+
+    def get_text_embeddings(
+        self,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        view_dependent_prompting: bool = True,
+    ) -> Float[Tensor, "BB N Nf"]:
+        batch_size = elevation.shape[0]
+
+        if view_dependent_prompting:
+            # Get direction
+            direction_idx = torch.zeros_like(elevation, dtype=torch.long)
+            for d in self.directions:
+                direction_idx[
+                    d.condition(elevation, azimuth, camera_distances)
+                ] = self.direction2idx[d.name]
+
+            # Get text embeddings
+            text_embeddings = self.text_embeddings_vd[direction_idx]  # type: ignore
+            uncond_text_embeddings = self.uncond_text_embeddings_vd[direction_idx]  # type: ignore
+        else:
+            text_embeddings = self.text_embeddings.expand(batch_size, -1, -1)  # type: ignore
+            uncond_text_embeddings = self.uncond_text_embeddings.expand(  # type: ignore
+                batch_size, -1, -1
+            )
+
+        # IMPORTANT: we return (cond, uncond), which is in different order than other implementations!
+        return torch.cat([text_embeddings, uncond_text_embeddings], dim=0)
+
+    def get_text_embeddings_perp_neg(
+        self,
+        elevation: Float[Tensor, "B"],
+        azimuth: Float[Tensor, "B"],
+        camera_distances: Float[Tensor, "B"],
+        view_dependent_prompting: bool = True,
+    ) -> Tuple[Float[Tensor, "BBBB N Nf"], Float[Tensor, "B 2"]]:
+        assert (
+            view_dependent_prompting
+        ), "Perp-Neg only works with view-dependent prompting"
+
+        batch_size = elevation.shape[0]
+
+        direction_idx = torch.zeros_like(elevation, dtype=torch.long)
+        for d in self.directions:
+            direction_idx[
+                d.condition(elevation, azimuth, camera_distances)
+            ] = self.direction2idx[d.name]
+        # 0 - side view
+        # 1 - front view
+        # 2 - back view
+        # 3 - overhead view
+
+        pos_text_embeddings = []
+        neg_text_embeddings = []
+        neg_guidance_weights = []
+        uncond_text_embeddings = []
+
+        side_emb = self.text_embeddings_vd[0]
+        front_emb = self.text_embeddings_vd[1]
+        back_emb = self.text_embeddings_vd[2]
+        overhead_emb = self.text_embeddings_vd[3]
+
+        for idx, ele, azi, dis in zip(
+            direction_idx, elevation, azimuth, camera_distances
+        ):
+            azi = shift_azimuth_deg(azi)  # to (-180, 180)
+            uncond_text_embeddings.append(
+                self.uncond_text_embeddings_vd[idx]
+            )  # should be ""
+            if idx.item() == 3:  # overhead view
+                pos_text_embeddings.append(overhead_emb)  # side view
+                # dummy
+                neg_text_embeddings += [
+                    self.uncond_text_embeddings_vd[idx],
+                    self.uncond_text_embeddings_vd[idx],
+                ]
+                neg_guidance_weights += [0.0, 0.0]
+            else:  # interpolating views
+                if torch.abs(azi) < 90:
+                    # front-side interpolation
+                    # 0 - complete side, 1 - complete front
+                    r_inter = 1 - torch.abs(azi) / 90
+                    pos_text_embeddings.append(
+                        r_inter * front_emb + (1 - r_inter) * side_emb
+                    )
+                    neg_text_embeddings += [front_emb, side_emb]
+                    neg_guidance_weights += [
+                        -shifted_expotional_decay(*self.perp_neg_f_fs, r_inter),
+                        -shifted_expotional_decay(*self.perp_neg_f_sf, 1 - r_inter),
+                    ]
+                else:
+                    # side-back interpolation
+                    # 0 - complete back, 1 - complete side
+                    r_inter = 2.0 - torch.abs(azi) / 90
+                    pos_text_embeddings.append(
+                        r_inter * side_emb + (1 - r_inter) * back_emb
+                    )
+                    neg_text_embeddings += [side_emb, front_emb]
+                    neg_guidance_weights += [
+                        -shifted_expotional_decay(*self.perp_neg_f_sb, r_inter),
+                        -shifted_expotional_decay(*self.perp_neg_f_fsb, r_inter),
+                    ]
+
+        text_embeddings = torch.cat(
+            [
+                torch.stack(pos_text_embeddings, dim=0),
+                torch.stack(uncond_text_embeddings, dim=0),
+                torch.stack(neg_text_embeddings, dim=0),
+            ],
+            dim=0,
+        )
+
+        return text_embeddings, torch.as_tensor(
+            neg_guidance_weights, device=elevation.device
+        ).reshape(batch_size, 2)
+
+
+def shift_azimuth_deg(azimuth: Float[Tensor, "..."]) -> Float[Tensor, "..."]:
+    # shift azimuth angle (in degrees), to [-180, 180]
+    return (azimuth + 180) % 360 - 180
+
+
+class PromptProcessor(BaseObject):
+    @dataclass
+    class Config(BaseObject.Config):
+        prompt: str = "a hamburger"
+
+        # manually assigned view-dependent prompts
+        prompt_front: Optional[str] = None
+        prompt_side: Optional[str] = None
+        prompt_back: Optional[str] = None
+        prompt_overhead: Optional[str] = None
+
+        negative_prompt: str = ""
+        pretrained_model_name_or_path: str = "runwayml/stable-diffusion-v1-5"
+        overhead_threshold: float = 60.0
+        front_threshold: float = 45.0
+        back_threshold: float = 45.0
+        view_dependent_prompt_front: bool = False
+        use_cache: bool = True
+        spawn: bool = True
+
+        # perp neg
+        use_perp_neg: bool = False
+        # a*e(-b*r) + c
+        # a * e(-b) + c = 0
+        perp_neg_f_sb: Tuple[float, float, float] = (1, 0.5, -0.606)
+        perp_neg_f_fsb: Tuple[float, float, float] = (1, 0.5, +0.967)
+        perp_neg_f_fs: Tuple[float, float, float] = (
+            4,
+            0.5,
+            -2.426,
+        )  # f_fs(1) = 0, a, b > 0
+        perp_neg_f_sf: Tuple[float, float, float] = (4, 0.5, -2.426)
+
+        # prompt debiasing
+        use_prompt_debiasing: bool = False
+        pretrained_model_name_or_path_prompt_debiasing: str = "bert-base-uncased"
+        # index of words that can potentially be removed
+        prompt_debiasing_mask_ids: Optional[List[int]] = None
+
+    cfg: Config
+
+    @rank_zero_only
+    def configure_text_encoder(self) -> None:
+        raise NotImplementedError
+
+    @rank_zero_only
+    def destroy_text_encoder(self) -> None:
+        raise NotImplementedError
+
+    def configure(self) -> None:
+        self._cache_dir = ".threestudio_cache/text_embeddings"  # FIXME: hard-coded path
+
+        # view-dependent text embeddings
+        self.directions: List[DirectionConfig]
+        if self.cfg.view_dependent_prompt_front:
+            self.directions = [
+                DirectionConfig(
+                    "side",
+                    lambda s: f"side view of {s}",
+                    lambda s: s,
+                    lambda ele, azi, dis: torch.ones_like(ele, dtype=torch.bool),
+                ),
+                DirectionConfig(
+                    "front",
+                    lambda s: f"front view of {s}",
+                    lambda s: s,
+                    lambda ele, azi, dis: (
+                        shift_azimuth_deg(azi) > -self.cfg.front_threshold
+                    )
+                    & (shift_azimuth_deg(azi) < self.cfg.front_threshold),
+                ),
+                DirectionConfig(
+                    "back",
+                    lambda s: f"backside view of {s}",
+                    lambda s: s,
+                    lambda ele, azi, dis: (
+                        shift_azimuth_deg(azi) > 180 - self.cfg.back_threshold
+                    )
+                    | (shift_azimuth_deg(azi) < -180 + self.cfg.back_threshold),
+                ),
+                DirectionConfig(
+                    "overhead",
+                    lambda s: f"overhead view of {s}",
+                    lambda s: s,
+                    lambda ele, azi, dis: ele > self.cfg.overhead_threshold,
+                ),
+            ]
+        else:
+            self.directions = [
+                DirectionConfig(
+                    "side",
+                    lambda s: f"{s}, side view",
+                    lambda s: s,
+                    lambda ele, azi, dis: torch.ones_like(ele, dtype=torch.bool),
+                ),
+                DirectionConfig(
+                    "front",
+                    lambda s: f"{s}, front view",
+                    lambda s: s,
+                    lambda ele, azi, dis: (
+                        shift_azimuth_deg(azi) > -self.cfg.front_threshold
+                    )
+                    & (shift_azimuth_deg(azi) < self.cfg.front_threshold),
+                ),
+                DirectionConfig(
+                    "back",
+                    lambda s: f"{s}, back view",
+                    lambda s: s,
+                    lambda ele, azi, dis: (
+                        shift_azimuth_deg(azi) > 180 - self.cfg.back_threshold
+                    )
+                    | (shift_azimuth_deg(azi) < -180 + self.cfg.back_threshold),
+                ),
+                DirectionConfig(
+                    "overhead",
+                    lambda s: f"{s}, overhead view",
+                    lambda s: s,
+                    lambda ele, azi, dis: ele > self.cfg.overhead_threshold,
+                ),
+            ]
+
+        self.direction2idx = {d.name: i for i, d in enumerate(self.directions)}
+
+        with open(os.path.join("load/prompt_library.json"), "r") as f:
+            self.prompt_library = json.load(f)
+        # use provided prompt or find prompt in library
+        self.prompt = self.preprocess_prompt(self.cfg.prompt)
+        # use provided negative prompt
+        self.negative_prompt = self.cfg.negative_prompt
+
+        threestudio.info(
+            f"Using prompt [{self.prompt}] and negative prompt [{self.negative_prompt}]"
+        )
+
+        # view-dependent prompting
+        if self.cfg.use_prompt_debiasing:
+            assert (
+                self.cfg.prompt_side is None
+                and self.cfg.prompt_back is None
+                and self.cfg.prompt_overhead is None
+            ), "Do not manually assign prompt_side, prompt_back or prompt_overhead when using prompt debiasing"
+            prompts = self.get_debiased_prompt(self.prompt)
+            self.prompts_vd = [
+                d.prompt(prompt) for d, prompt in zip(self.directions, prompts)
+            ]
+        else:
+            self.prompts_vd = [
+                self.cfg.get(f"prompt_{d.name}", None) or d.prompt(self.prompt)  # type: ignore
+                for d in self.directions
+            ]
+
+        prompts_vd_display = " ".join(
+            [
+                f"[{d.name}]:[{prompt}]"
+                for prompt, d in zip(self.prompts_vd, self.directions)
+            ]
+        )
+        threestudio.info(f"Using view-dependent prompts {prompts_vd_display}")
+
+        self.negative_prompts_vd = [
+            d.negative_prompt(self.negative_prompt) for d in self.directions
+        ]
+
+        self.prepare_text_embeddings()
+        self.load_text_embeddings()
+
+    @staticmethod
+    def spawn_func(pretrained_model_name_or_path, prompts, cache_dir, device):
+        raise NotImplementedError
+
+    @rank_zero_only
+    def prepare_text_embeddings(self):
+        os.makedirs(self._cache_dir, exist_ok=True)
+
+        all_prompts = (
+            [self.prompt]
+            + [self.negative_prompt]
+            + self.prompts_vd
+            + self.negative_prompts_vd
+        )
+        prompts_to_process = []
+        for prompt in all_prompts:
+            if self.cfg.use_cache:
+                # some text embeddings are already in cache
+                # do not process them
+                cache_path = os.path.join(
+                    self._cache_dir,
+                    f"{hash_prompt(self.cfg.pretrained_model_name_or_path, prompt)}.pt",
+                )
+                if os.path.exists(cache_path):
+                    threestudio.debug(
+                        f"Text embeddings for model {self.cfg.pretrained_model_name_or_path} and prompt [{prompt}] are already in cache, skip processing."
+                    )
+                    continue
+            prompts_to_process.append(prompt)
+
+        if len(prompts_to_process) > 0:
+            if self.cfg.spawn:
+                ctx = mp.get_context("spawn")
+                subprocess = ctx.Process(
+                    target=self.spawn_func,
+                    args=(
+                        self.cfg.pretrained_model_name_or_path,
+                        prompts_to_process,
+                        self._cache_dir,
+                        self.device
+                    ),
+                )
+                subprocess.start()
+                subprocess.join()
+            else:
+                self.spawn_func(
+                    self.cfg.pretrained_model_name_or_path,
+                    prompts_to_process,
+                    self._cache_dir,
+                    self.device
+                )
+            cleanup()
+
+    def load_text_embeddings(self):
+        # synchronize, to ensure the text embeddings have been computed and saved to cache
+        barrier()
+        self.text_embeddings = self.load_from_cache(self.prompt)[None, ...]
+        self.uncond_text_embeddings = self.load_from_cache(self.negative_prompt)[
+            None, ...
+        ]
+        self.text_embeddings_vd = torch.stack(
+            [self.load_from_cache(prompt) for prompt in self.prompts_vd], dim=0
+        )
+        self.uncond_text_embeddings_vd = torch.stack(
+            [self.load_from_cache(prompt) for prompt in self.negative_prompts_vd], dim=0
+        )
+        threestudio.debug(f"Loaded text embeddings.")
+
+    def load_from_cache(self, prompt):
+        cache_path = os.path.join(
+            self._cache_dir,
+            f"{hash_prompt(self.cfg.pretrained_model_name_or_path, prompt)}.pt",
+        )
+        if not os.path.exists(cache_path):
+            raise FileNotFoundError(
+                f"Text embedding file {cache_path} for model {self.cfg.pretrained_model_name_or_path} and prompt [{prompt}] not found."
+            )
+        return torch.load(cache_path, map_location=self.device)
+
+    def preprocess_prompt(self, prompt: str) -> str:
+        if prompt.startswith("lib:"):
+            # find matches in the library
+            candidate = None
+            keywords = prompt[4:].lower().split("_")
+            for prompt in self.prompt_library["dreamfusion"]:
+                if all([k in prompt.lower() for k in keywords]):
+                    if candidate is not None:
+                        raise ValueError(
+                            f"Multiple prompts matched with keywords {keywords} in library"
+                        )
+                    candidate = prompt
+            if candidate is None:
+                raise ValueError(
+                    f"Cannot find prompt with keywords {keywords} in library"
+                )
+            threestudio.info("Find matched prompt in library: " + candidate)
+            return candidate
+        else:
+            return prompt
+
+    def get_text_embeddings(
+        self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]]
+    ) -> Tuple[Float[Tensor, "B ..."], Float[Tensor, "B ..."]]:
+        raise NotImplementedError
+
+    def get_debiased_prompt(self, prompt: str) -> List[str]:
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            self.cfg.pretrained_model_name_or_path_prompt_debiasing
+        )
+        model = BertForMaskedLM.from_pretrained(
+            self.cfg.pretrained_model_name_or_path_prompt_debiasing
+        )
+
+        views = [d.name for d in self.directions]
+        view_ids = tokenizer(" ".join(views), return_tensors="pt").input_ids[0]
+        view_ids = view_ids[1:5]
+
+        def modulate(prompt):
+            prompt_vd = f"This image is depicting a [MASK] view of {prompt}"
+            tokens = tokenizer(
+                prompt_vd,
+                padding="max_length",
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            mask_idx = torch.where(tokens.input_ids == tokenizer.mask_token_id)[1]
+
+            logits = model(**tokens).logits
+            logits = F.softmax(logits[0, mask_idx], dim=-1)
+            logits = logits[0, view_ids]
+            probes = logits / logits.sum()
+            return probes
+
+        prompts = [prompt.split(" ") for _ in range(4)]
+        full_probe = modulate(prompt)
+        n_words = len(prompt.split(" "))
+        prompt_debiasing_mask_ids = (
+            self.cfg.prompt_debiasing_mask_ids
+            if self.cfg.prompt_debiasing_mask_ids is not None
+            else list(range(n_words))
+        )
+        words_to_debias = [prompt.split(" ")[idx] for idx in prompt_debiasing_mask_ids]
+        threestudio.info(f"Words that can potentially be removed: {words_to_debias}")
+        for idx in prompt_debiasing_mask_ids:
+            words = prompt.split(" ")
+            prompt_ = " ".join(words[:idx] + words[(idx + 1) :])
+            part_probe = modulate(prompt_)
+
+            pmi = full_probe / torch.lerp(part_probe, full_probe, 0.5)
+            for i in range(pmi.shape[0]):
+                if pmi[i].item() < 0.95:
+                    prompts[i][idx] = ""
+
+        debiased_prompts = [" ".join([word for word in p if word]) for p in prompts]
+        for d, debiased_prompt in zip(views, debiased_prompts):
+            threestudio.info(f"Debiased prompt of the {d} view is [{debiased_prompt}]")
+
+        del tokenizer, model
+        cleanup()
+
+        return debiased_prompts
+
+    def __call__(self) -> PromptProcessorOutput:
+        return PromptProcessorOutput(
+            text_embeddings=self.text_embeddings,
+            uncond_text_embeddings=self.uncond_text_embeddings,
+            text_embeddings_vd=self.text_embeddings_vd,
+            uncond_text_embeddings_vd=self.uncond_text_embeddings_vd,
+            directions=self.directions,
+            direction2idx=self.direction2idx,
+            use_perp_neg=self.cfg.use_perp_neg,
+            perp_neg_f_sb=self.cfg.perp_neg_f_sb,
+            perp_neg_f_fsb=self.cfg.perp_neg_f_fsb,
+            perp_neg_f_fs=self.cfg.perp_neg_f_fs,
+            perp_neg_f_sf=self.cfg.perp_neg_f_sf,
+        )

threestudio/models/prompt_processors/clip_prompt_processor.py

@@ -0,0 +1,44 @@
+import json
+import os
+from dataclasses import dataclass
+
+import clip
+import torch
+import torch
+import torch.nn as nn
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessor, hash_prompt
+from threestudio.utils.misc import cleanup
+from threestudio.utils.typing import *
+
+
+@threestudio.register("clip-prompt-processor")
+class ClipPromptProcessor(PromptProcessor):
+    @dataclass
+    class Config(PromptProcessor.Config):
+        pass
+
+    cfg: Config
+
+    @staticmethod
+    def spawn_func(pretrained_model_name_or_path, prompts, cache_dir, device):
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        clip_model, _ = clip.load(pretrained_model_name_or_path, jit=False)
+        with torch.no_grad():
+            tokens = clip.tokenize(
+                prompts,
+            ).to(device)
+            text_embeddings = clip_model.encode_text(tokens)
+            text_embeddings = text_embeddings / text_embeddings.norm(dim=-1, keepdim=True)
+
+        for prompt, embedding in zip(prompts, text_embeddings):
+            torch.save(
+                embedding,
+                os.path.join(
+                    cache_dir,
+                    f"{hash_prompt(pretrained_model_name_or_path, prompt)}.pt",
+                ),
+            )
+
+        del clip_model

threestudio/models/prompt_processors/deepfloyd_prompt_processor.py

@@ -0,0 +1,98 @@
+import json
+import os
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from diffusers import IFPipeline
+from transformers import T5EncoderModel, T5Tokenizer
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessor, hash_prompt
+from threestudio.utils.misc import cleanup
+from threestudio.utils.typing import *
+
+
+@threestudio.register("deep-floyd-prompt-processor")
+class DeepFloydPromptProcessor(PromptProcessor):
+    @dataclass
+    class Config(PromptProcessor.Config):
+        pretrained_model_name_or_path: str = "DeepFloyd/IF-I-XL-v1.0"
+
+    cfg: Config
+
+    ### these functions are unused, kept for debugging ###
+    def configure_text_encoder(self) -> None:
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        self.text_encoder = T5EncoderModel.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            subfolder="text_encoder",
+            load_in_8bit=True,
+            variant="8bit",
+            device_map="auto",
+        )  # FIXME: behavior of auto device map in multi-GPU training
+        self.pipe = IFPipeline.from_pretrained(
+            self.cfg.pretrained_model_name_or_path,
+            text_encoder=self.text_encoder,  # pass the previously instantiated 8bit text encoder
+            unet=None,
+        )
+
+    def destroy_text_encoder(self) -> None:
+        del self.text_encoder
+        del self.pipe
+        cleanup()
+
+    def get_text_embeddings(
+        self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]]
+    ) -> Tuple[Float[Tensor, "B 77 4096"], Float[Tensor, "B 77 4096"]]:
+        text_embeddings, uncond_text_embeddings = self.pipe.encode_prompt(
+            prompt=prompt, negative_prompt=negative_prompt, device=self.device
+        )
+        return text_embeddings, uncond_text_embeddings
+
+    ###
+
+    @staticmethod
+    def spawn_func(pretrained_model_name_or_path, prompts, cache_dir, device):
+        max_length = 77
+        tokenizer = T5Tokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="tokenizer",
+            local_files_only=True
+        )
+        text_encoder = T5EncoderModel.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="text_encoder",
+            torch_dtype=torch.float16,  # suppress warning
+            load_in_8bit=True,
+            variant="8bit",
+            device_map="auto",
+            local_files_only=True
+        )
+        with torch.no_grad():
+            text_inputs = tokenizer(
+                prompts,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            attention_mask = text_inputs.attention_mask
+            text_embeddings = text_encoder(
+                text_input_ids.to(text_encoder.device),
+                attention_mask=attention_mask.to(text_encoder.device),
+            )
+            text_embeddings = text_embeddings[0]
+
+        for prompt, embedding in zip(prompts, text_embeddings):
+            torch.save(
+                embedding,
+                os.path.join(
+                    cache_dir,
+                    f"{hash_prompt(pretrained_model_name_or_path, prompt)}.pt",
+                ),
+            )
+
+        del text_encoder

threestudio/models/prompt_processors/dummy_prompt_processor.py

@@ -0,0 +1,18 @@
+import json
+import os
+from dataclasses import dataclass
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessor, hash_prompt
+from threestudio.utils.misc import cleanup
+from threestudio.utils.typing import *
+
+
+@threestudio.register("dummy-prompt-processor")
+class DummyPromptProcessor(PromptProcessor):
+    @dataclass
+    class Config(PromptProcessor.Config):
+        pretrained_model_name_or_path: str = ""
+        prompt: str = ""
+
+    cfg: Config

threestudio/models/prompt_processors/stable_diffusion_prompt_processor.py

@@ -0,0 +1,136 @@
+import json
+import os
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, CLIPTextModel
+
+import threestudio
+from threestudio.models.prompt_processors.base import PromptProcessor, hash_prompt
+from threestudio.utils.misc import cleanup
+from threestudio.utils.typing import *
+
+
+@threestudio.register("stable-diffusion-prompt-processor")
+class StableDiffusionPromptProcessor(PromptProcessor):
+    @dataclass
+    class Config(PromptProcessor.Config):
+        pass
+
+    cfg: Config
+
+    ### these functions are unused, kept for debugging ###
+    def configure_text_encoder(self) -> None:
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.cfg.pretrained_model_name_or_path, subfolder="tokenizer"
+        )
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        self.text_encoder = CLIPTextModel.from_pretrained(
+            self.cfg.pretrained_model_name_or_path, subfolder="text_encoder"
+        ).to(self.device)
+
+        for p in self.text_encoder.parameters():
+            p.requires_grad_(False)
+
+    def destroy_text_encoder(self) -> None:
+        del self.tokenizer
+        del self.text_encoder
+        cleanup()
+
+    def get_text_embeddings(
+        self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]]
+    ) -> Tuple[Float[Tensor, "B 77 768"], Float[Tensor, "B 77 768"]]:
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if isinstance(negative_prompt, str):
+            negative_prompt = [negative_prompt]
+        # Tokenize text and get embeddings
+        tokens = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        )
+        uncond_tokens = self.tokenizer(
+            negative_prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        )
+
+        with torch.no_grad():
+            text_embeddings = self.text_encoder(tokens.input_ids.to(self.device))[0]
+            uncond_text_embeddings = self.text_encoder(
+                uncond_tokens.input_ids.to(self.device)
+            )[0]
+
+        return text_embeddings, uncond_text_embeddings
+
+    ###
+
+    @staticmethod
+    def spawn_func(pretrained_model_name_or_path, prompts, cache_dir, device):
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="tokenizer",
+            local_files_only=True,
+        )
+        text_encoder = CLIPTextModel.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="text_encoder",
+            device_map="auto",
+            local_files_only=True,
+        )
+
+        with torch.no_grad():
+            tokens = tokenizer(
+                prompts,
+                padding="max_length",
+                max_length=tokenizer.model_max_length,
+                return_tensors="pt",
+            )
+            text_embeddings = text_encoder(tokens.input_ids.to(text_encoder.device))[0]
+
+        for prompt, embedding in zip(prompts, text_embeddings):
+            torch.save(
+                embedding,
+                os.path.join(
+                    cache_dir,
+                    f"{hash_prompt(pretrained_model_name_or_path, prompt)}.pt",
+                ),
+            )
+
+        del text_encoder
+
+
+from transformers.models.clip import CLIPTextModel, CLIPTokenizer
+def add_tokens_to_model(learned_embeds_path, text_encoder: CLIPTextModel, 
+        tokenizer: CLIPTokenizer, override_token: Optional[Union[str, dict]] = None) -> None:
+    r"""Adds tokens to the tokenizer and text encoder of a model."""
+
+    learned_embeds = torch.load(learned_embeds_path, map_location='cpu')
+    
+    # Loop over learned embeddings
+    new_tokens = []
+    for token, embedding in learned_embeds.items():
+        embedding = embedding.to(text_encoder.get_input_embeddings().weight.dtype)
+        if override_token is not None:
+            token = override_token if isinstance(override_token, str) else override_token[token]
+        
+        # Add the token to the tokenizer
+        num_added_tokens = tokenizer.add_tokens(token)
+        if num_added_tokens == 0:
+            raise ValueError((f"The tokenizer already contains the token {token}. Please pass a "
+                               "different `token` that is not already in the tokenizer."))
+  
+        # Resize the token embeddings
+        text_encoder.resize_token_embeddings(len(tokenizer))
+  
+        # Get the id for the token and assign the embeds
+        token_id = tokenizer.convert_tokens_to_ids(token)
+        text_encoder.get_input_embeddings().weight.data[token_id] = embedding    
+        new_tokens.append(token)
+
+    print(f'Added {len(new_tokens)} tokens to tokenizer and text embedding: {new_tokens}')

threestudio/models/renderers/__init__.py

@@ -0,0 +1,9 @@
+from . import (
+    base,
+    deferred_volume_renderer,
+    gan_volume_renderer,
+    nerf_volume_renderer,
+    neus_volume_renderer,
+    nvdiff_rasterizer,
+    patch_renderer,
+)

threestudio/models/renderers/base.py

@@ -0,0 +1,80 @@
+from dataclasses import dataclass
+
+import nerfacc
+import torch
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.geometry.base import BaseImplicitGeometry
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.utils.base import BaseModule
+from threestudio.utils.typing import *
+
+
+class Renderer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        radius: float = 1.0
+
+    cfg: Config
+
+    def configure(
+        self,
+        geometry: BaseImplicitGeometry,
+        material: BaseMaterial,
+        background: BaseBackground,
+    ) -> None:
+        # keep references to submodules using namedtuple, avoid being registered as modules
+        @dataclass
+        class SubModules:
+            geometry: BaseImplicitGeometry
+            material: BaseMaterial
+            background: BaseBackground
+
+        self.sub_modules = SubModules(geometry, material, background)
+
+        # set up bounding box
+        self.bbox: Float[Tensor, "2 3"]
+        self.register_buffer(
+            "bbox",
+            torch.as_tensor(
+                [
+                    [-self.cfg.radius, -self.cfg.radius, -self.cfg.radius],
+                    [self.cfg.radius, self.cfg.radius, self.cfg.radius],
+                ],
+                dtype=torch.float32,
+            ),
+        )
+
+    def forward(self, *args, **kwargs) -> Dict[str, Any]:
+        raise NotImplementedError
+
+    @property
+    def geometry(self) -> BaseImplicitGeometry:
+        return self.sub_modules.geometry
+
+    @property
+    def material(self) -> BaseMaterial:
+        return self.sub_modules.material
+
+    @property
+    def background(self) -> BaseBackground:
+        return self.sub_modules.background
+
+    def set_geometry(self, geometry: BaseImplicitGeometry) -> None:
+        self.sub_modules.geometry = geometry
+
+    def set_material(self, material: BaseMaterial) -> None:
+        self.sub_modules.material = material
+
+    def set_background(self, background: BaseBackground) -> None:
+        self.sub_modules.background = background
+
+
+class VolumeRenderer(Renderer):
+    pass
+
+
+class Rasterizer(Renderer):
+    pass

threestudio/models/renderers/deferred_volume_renderer.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.renderers.base import VolumeRenderer
+
+
+class DeferredVolumeRenderer(VolumeRenderer):
+    pass

threestudio/models/renderers/gan_volume_renderer.py

@@ -0,0 +1,159 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.geometry.base import BaseImplicitGeometry
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.renderers.base import VolumeRenderer
+from threestudio.utils.GAN.discriminator import NLayerDiscriminator, weights_init
+from threestudio.utils.GAN.distribution import DiagonalGaussianDistribution
+from threestudio.utils.GAN.mobilenet import MobileNetV3 as GlobalEncoder
+from threestudio.utils.GAN.vae import Decoder as Generator
+from threestudio.utils.GAN.vae import Encoder as LocalEncoder
+from threestudio.utils.typing import *
+
+
+@threestudio.register("gan-volume-renderer")
+class GANVolumeRenderer(VolumeRenderer):
+    @dataclass
+    class Config(VolumeRenderer.Config):
+        base_renderer_type: str = ""
+        base_renderer: Optional[VolumeRenderer.Config] = None
+
+    cfg: Config
+
+    def configure(
+        self,
+        geometry: BaseImplicitGeometry,
+        material: BaseMaterial,
+        background: BaseBackground,
+    ) -> None:
+        self.base_renderer = threestudio.find(self.cfg.base_renderer_type)(
+            self.cfg.base_renderer,
+            geometry=geometry,
+            material=material,
+            background=background,
+        )
+        self.ch_mult = [1, 2, 4]
+        self.generator = Generator(
+            ch=64,
+            out_ch=3,
+            ch_mult=self.ch_mult,
+            num_res_blocks=1,
+            attn_resolutions=[],
+            dropout=0.0,
+            resamp_with_conv=True,
+            in_channels=7,
+            resolution=512,
+            z_channels=4,
+        )
+        self.local_encoder = LocalEncoder(
+            ch=32,
+            out_ch=3,
+            ch_mult=self.ch_mult,
+            num_res_blocks=1,
+            attn_resolutions=[],
+            dropout=0.0,
+            resamp_with_conv=True,
+            in_channels=3,
+            resolution=512,
+            z_channels=4,
+        )
+        self.global_encoder = GlobalEncoder(n_class=64)
+        self.discriminator = NLayerDiscriminator(
+            input_nc=3, n_layers=3, use_actnorm=False, ndf=64
+        ).apply(weights_init)
+
+    def forward(
+        self,
+        rays_o: Float[Tensor, "B H W 3"],
+        rays_d: Float[Tensor, "B H W 3"],
+        light_positions: Float[Tensor, "B 3"],
+        bg_color: Optional[Tensor] = None,
+        gt_rgb: Float[Tensor, "B H W 3"] = None,
+        multi_level_guidance: Bool = False,
+        **kwargs
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        B, H, W, _ = rays_o.shape
+        if gt_rgb is not None and multi_level_guidance:
+            generator_level = torch.randint(0, 3, (1,)).item()
+            interval_x = torch.randint(0, 8, (1,)).item()
+            interval_y = torch.randint(0, 8, (1,)).item()
+            int_rays_o = rays_o[:, interval_y::8, interval_x::8]
+            int_rays_d = rays_d[:, interval_y::8, interval_x::8]
+            out = self.base_renderer(
+                int_rays_o, int_rays_d, light_positions, bg_color, **kwargs
+            )
+            comp_int_rgb = out["comp_rgb"][..., :3]
+            comp_gt_rgb = gt_rgb[:, interval_y::8, interval_x::8]
+        else:
+            generator_level = 0
+        scale_ratio = 2 ** (len(self.ch_mult) - 1)
+        rays_o = torch.nn.functional.interpolate(
+            rays_o.permute(0, 3, 1, 2),
+            (H // scale_ratio, W // scale_ratio),
+            mode="bilinear",
+        ).permute(0, 2, 3, 1)
+        rays_d = torch.nn.functional.interpolate(
+            rays_d.permute(0, 3, 1, 2),
+            (H // scale_ratio, W // scale_ratio),
+            mode="bilinear",
+        ).permute(0, 2, 3, 1)
+
+        out = self.base_renderer(rays_o, rays_d, light_positions, bg_color, **kwargs)
+        comp_rgb = out["comp_rgb"][..., :3]
+        latent = out["comp_rgb"][..., 3:]
+        out["comp_lr_rgb"] = comp_rgb.clone()
+
+        posterior = DiagonalGaussianDistribution(latent.permute(0, 3, 1, 2))
+        if multi_level_guidance:
+            z_map = posterior.sample()
+        else:
+            z_map = posterior.mode()
+        lr_rgb = comp_rgb.permute(0, 3, 1, 2)
+
+        if generator_level == 0:
+            g_code_rgb = self.global_encoder(F.interpolate(lr_rgb, (224, 224)))
+            comp_gan_rgb = self.generator(torch.cat([lr_rgb, z_map], dim=1), g_code_rgb)
+        elif generator_level == 1:
+            g_code_rgb = self.global_encoder(
+                F.interpolate(gt_rgb.permute(0, 3, 1, 2), (224, 224))
+            )
+            comp_gan_rgb = self.generator(torch.cat([lr_rgb, z_map], dim=1), g_code_rgb)
+        elif generator_level == 2:
+            g_code_rgb = self.global_encoder(
+                F.interpolate(gt_rgb.permute(0, 3, 1, 2), (224, 224))
+            )
+            l_code_rgb = self.local_encoder(gt_rgb.permute(0, 3, 1, 2))
+            posterior = DiagonalGaussianDistribution(l_code_rgb)
+            z_map = posterior.sample()
+            comp_gan_rgb = self.generator(torch.cat([lr_rgb, z_map], dim=1), g_code_rgb)
+
+        comp_rgb = F.interpolate(comp_rgb.permute(0, 3, 1, 2), (H, W), mode="bilinear")
+        comp_gan_rgb = F.interpolate(comp_gan_rgb, (H, W), mode="bilinear")
+        out.update(
+            {
+                "posterior": posterior,
+                "comp_gan_rgb": comp_gan_rgb.permute(0, 2, 3, 1),
+                "comp_rgb": comp_rgb.permute(0, 2, 3, 1),
+                "generator_level": generator_level,
+            }
+        )
+
+        if gt_rgb is not None and multi_level_guidance:
+            out.update({"comp_int_rgb": comp_int_rgb, "comp_gt_rgb": comp_gt_rgb})
+        return out
+
+    def update_step(
+        self, epoch: int, global_step: int, on_load_weights: bool = False
+    ) -> None:
+        self.base_renderer.update_step(epoch, global_step, on_load_weights)
+
+    def train(self, mode=True):
+        return self.base_renderer.train(mode)
+
+    def eval(self):
+        return self.base_renderer.eval()

threestudio/models/renderers/nerf_volume_renderer.py

@@ -0,0 +1,462 @@
+from dataclasses import dataclass, field
+from functools import partial
+
+import nerfacc
+import torch
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.estimators import ImportanceEstimator
+from threestudio.models.geometry.base import BaseImplicitGeometry
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.networks import create_network_with_input_encoding
+from threestudio.models.renderers.base import VolumeRenderer
+from threestudio.systems.utils import parse_optimizer, parse_scheduler_to_instance
+from threestudio.utils.ops import chunk_batch, get_activation, validate_empty_rays
+from threestudio.utils.typing import *
+
+
+@threestudio.register("nerf-volume-renderer")
+class NeRFVolumeRenderer(VolumeRenderer):
+    @dataclass
+    class Config(VolumeRenderer.Config):
+        num_samples_per_ray: int = 512
+        eval_chunk_size: int = 160000
+        randomized: bool = True
+
+        near_plane: float = 0.0
+        far_plane: float = 1e10
+
+        return_comp_normal: bool = False
+        return_normal_perturb: bool = False
+
+        # in ["occgrid", "proposal", "importance"]
+        estimator: str = "occgrid"
+
+        # for occgrid
+        grid_prune: bool = True
+        prune_alpha_threshold: bool = True
+
+        # for proposal
+        proposal_network_config: Optional[dict] = None
+        prop_optimizer_config: Optional[dict] = None
+        prop_scheduler_config: Optional[dict] = None
+        num_samples_per_ray_proposal: int = 64
+
+        # for importance
+        num_samples_per_ray_importance: int = 64
+
+    cfg: Config
+
+    def configure(
+        self,
+        geometry: BaseImplicitGeometry,
+        material: BaseMaterial,
+        background: BaseBackground,
+    ) -> None:
+        super().configure(geometry, material, background)
+        if self.cfg.estimator == "occgrid":
+            self.estimator = nerfacc.OccGridEstimator(
+                roi_aabb=self.bbox.view(-1), resolution=32, levels=1
+            )
+            if not self.cfg.grid_prune:
+                self.estimator.occs.fill_(True)
+                self.estimator.binaries.fill_(True)
+            self.render_step_size = (
+                1.732 * 2 * self.cfg.radius / self.cfg.num_samples_per_ray
+            )
+            self.randomized = self.cfg.randomized
+        elif self.cfg.estimator == "importance":
+            self.estimator = ImportanceEstimator()
+        elif self.cfg.estimator == "proposal":
+            self.prop_net = create_network_with_input_encoding(
+                **self.cfg.proposal_network_config
+            )
+            self.prop_optim = parse_optimizer(
+                self.cfg.prop_optimizer_config, self.prop_net
+            )
+            self.prop_scheduler = (
+                parse_scheduler_to_instance(
+                    self.cfg.prop_scheduler_config, self.prop_optim
+                )
+                if self.cfg.prop_scheduler_config is not None
+                else None
+            )
+            self.estimator = nerfacc.PropNetEstimator(
+                self.prop_optim, self.prop_scheduler
+            )
+
+            def get_proposal_requires_grad_fn(
+                target: float = 5.0, num_steps: int = 1000
+            ):
+                schedule = lambda s: min(s / num_steps, 1.0) * target
+
+                steps_since_last_grad = 0
+
+                def proposal_requires_grad_fn(step: int) -> bool:
+                    nonlocal steps_since_last_grad
+                    target_steps_since_last_grad = schedule(step)
+                    requires_grad = steps_since_last_grad > target_steps_since_last_grad
+                    if requires_grad:
+                        steps_since_last_grad = 0
+                    steps_since_last_grad += 1
+                    return requires_grad
+
+                return proposal_requires_grad_fn
+
+            self.proposal_requires_grad_fn = get_proposal_requires_grad_fn()
+            self.randomized = self.cfg.randomized
+        else:
+            raise NotImplementedError(
+                "Unknown estimator, should be one of ['occgrid', 'proposal', 'importance']."
+            )
+
+        # for proposal
+        self.vars_in_forward = {}
+
+    def forward(
+        self,
+        rays_o: Float[Tensor, "B H W 3"],
+        rays_d: Float[Tensor, "B H W 3"],
+        light_positions: Float[Tensor, "B 3"],
+        bg_color: Optional[Tensor] = None,
+        **kwargs
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        batch_size, height, width = rays_o.shape[:3]
+        rays_o_flatten: Float[Tensor, "Nr 3"] = rays_o.reshape(-1, 3)
+        rays_d_flatten: Float[Tensor, "Nr 3"] = rays_d.reshape(-1, 3)
+        light_positions_flatten: Float[Tensor, "Nr 3"] = (
+            light_positions.reshape(-1, 1, 1, 3)
+            .expand(-1, height, width, -1)
+            .reshape(-1, 3)
+        )
+        n_rays = rays_o_flatten.shape[0]
+
+        if self.cfg.estimator == "occgrid":
+            if not self.cfg.grid_prune:
+                with torch.no_grad():
+                    ray_indices, t_starts_, t_ends_ = self.estimator.sampling(
+                        rays_o_flatten,
+                        rays_d_flatten,
+                        sigma_fn=None,
+                        near_plane=self.cfg.near_plane,
+                        far_plane=self.cfg.far_plane,
+                        render_step_size=self.render_step_size,
+                        alpha_thre=0.0,
+                        stratified=self.randomized,
+                        cone_angle=0.0,
+                        early_stop_eps=0,
+                    )
+            else:
+
+                def sigma_fn(t_starts, t_ends, ray_indices):
+                    t_starts, t_ends = t_starts[..., None], t_ends[..., None]
+                    t_origins = rays_o_flatten[ray_indices]
+                    t_positions = (t_starts + t_ends) / 2.0
+                    t_dirs = rays_d_flatten[ray_indices]
+                    positions = t_origins + t_dirs * t_positions
+                    if self.training:
+                        sigma = self.geometry.forward_density(positions)[..., 0]
+                    else:
+                        sigma = chunk_batch(
+                            self.geometry.forward_density,
+                            self.cfg.eval_chunk_size,
+                            positions,
+                        )[..., 0]
+                    return sigma
+
+                with torch.no_grad():
+                    ray_indices, t_starts_, t_ends_ = self.estimator.sampling(
+                        rays_o_flatten,
+                        rays_d_flatten,
+                        sigma_fn=sigma_fn if self.cfg.prune_alpha_threshold else None,
+                        near_plane=self.cfg.near_plane,
+                        far_plane=self.cfg.far_plane,
+                        render_step_size=self.render_step_size,
+                        alpha_thre=0.01 if self.cfg.prune_alpha_threshold else 0.0,
+                        stratified=self.randomized,
+                        cone_angle=0.0,
+                    )
+        elif self.cfg.estimator == "proposal":
+
+            def prop_sigma_fn(
+                t_starts: Float[Tensor, "Nr Ns"],
+                t_ends: Float[Tensor, "Nr Ns"],
+                proposal_network,
+            ):
+                t_origins: Float[Tensor, "Nr 1 3"] = rays_o_flatten.unsqueeze(-2)
+                t_dirs: Float[Tensor, "Nr 1 3"] = rays_d_flatten.unsqueeze(-2)
+                positions: Float[Tensor, "Nr Ns 3"] = (
+                    t_origins + t_dirs * (t_starts + t_ends)[..., None] / 2.0
+                )
+                aabb_min, aabb_max = self.bbox[0], self.bbox[1]
+                positions = (positions - aabb_min) / (aabb_max - aabb_min)
+                selector = ((positions > 0.0) & (positions < 1.0)).all(dim=-1)
+                density_before_activation = (
+                    proposal_network(positions.view(-1, 3))
+                    .view(*positions.shape[:-1], 1)
+                    .to(positions)
+                )
+                density: Float[Tensor, "Nr Ns 1"] = (
+                    get_activation("shifted_trunc_exp")(density_before_activation)
+                    * selector[..., None]
+                )
+                return density.squeeze(-1)
+
+            t_starts_, t_ends_ = self.estimator.sampling(
+                prop_sigma_fns=[partial(prop_sigma_fn, proposal_network=self.prop_net)],
+                prop_samples=[self.cfg.num_samples_per_ray_proposal],
+                num_samples=self.cfg.num_samples_per_ray,
+                n_rays=n_rays,
+                near_plane=self.cfg.near_plane,
+                far_plane=self.cfg.far_plane,
+                sampling_type="uniform",
+                stratified=self.randomized,
+                requires_grad=self.vars_in_forward["requires_grad"],
+            )
+            ray_indices = (
+                torch.arange(n_rays, device=rays_o_flatten.device)
+                .unsqueeze(-1)
+                .expand(-1, t_starts_.shape[1])
+            )
+            ray_indices = ray_indices.flatten()
+            t_starts_ = t_starts_.flatten()
+            t_ends_ = t_ends_.flatten()
+        elif self.cfg.estimator == "importance":
+
+            def prop_sigma_fn(
+                t_starts: Float[Tensor, "Nr Ns"],
+                t_ends: Float[Tensor, "Nr Ns"],
+                proposal_network,
+            ):
+                t_origins: Float[Tensor, "Nr 1 3"] = rays_o_flatten.unsqueeze(-2)
+                t_dirs: Float[Tensor, "Nr 1 3"] = rays_d_flatten.unsqueeze(-2)
+                positions: Float[Tensor, "Nr Ns 3"] = (
+                    t_origins + t_dirs * (t_starts + t_ends)[..., None] / 2.0
+                )
+                with torch.no_grad():
+                    geo_out = chunk_batch(
+                        proposal_network,
+                        self.cfg.eval_chunk_size,
+                        positions.reshape(-1, 3),
+                        output_normal=False,
+                    )
+                    density = geo_out["density"]
+                return density.reshape(positions.shape[:2])
+
+            t_starts_, t_ends_ = self.estimator.sampling(
+                prop_sigma_fns=[partial(prop_sigma_fn, proposal_network=self.geometry)],
+                prop_samples=[self.cfg.num_samples_per_ray_importance],
+                num_samples=self.cfg.num_samples_per_ray,
+                n_rays=n_rays,
+                near_plane=self.cfg.near_plane,
+                far_plane=self.cfg.far_plane,
+                sampling_type="uniform",
+                stratified=self.randomized,
+            )
+            ray_indices = (
+                torch.arange(n_rays, device=rays_o_flatten.device)
+                .unsqueeze(-1)
+                .expand(-1, t_starts_.shape[1])
+            )
+            ray_indices = ray_indices.flatten()
+            t_starts_ = t_starts_.flatten()
+            t_ends_ = t_ends_.flatten()
+        else:
+            raise NotImplementedError
+
+        ray_indices, t_starts_, t_ends_ = validate_empty_rays(
+            ray_indices, t_starts_, t_ends_
+        )
+        ray_indices = ray_indices.long()
+        t_starts, t_ends = t_starts_[..., None], t_ends_[..., None]
+        t_origins = rays_o_flatten[ray_indices]
+        t_dirs = rays_d_flatten[ray_indices]
+        t_light_positions = light_positions_flatten[ray_indices]
+        t_positions = (t_starts + t_ends) / 2.0
+        positions = t_origins + t_dirs * t_positions
+        t_intervals = t_ends - t_starts
+
+        if self.training:
+            geo_out = self.geometry(
+                positions, output_normal=self.material.requires_normal
+            )
+            rgb_fg_all = self.material(
+                viewdirs=t_dirs,
+                positions=positions,
+                light_positions=t_light_positions,
+                **geo_out,
+                **kwargs
+            )
+            comp_rgb_bg = self.background(dirs=rays_d)
+        else:
+            geo_out = chunk_batch(
+                self.geometry,
+                self.cfg.eval_chunk_size,
+                positions,
+                output_normal=self.material.requires_normal,
+            )
+            rgb_fg_all = chunk_batch(
+                self.material,
+                self.cfg.eval_chunk_size,
+                viewdirs=t_dirs,
+                positions=positions,
+                light_positions=t_light_positions,
+                **geo_out
+            )
+            comp_rgb_bg = chunk_batch(
+                self.background, self.cfg.eval_chunk_size, dirs=rays_d
+            )
+
+        weights: Float[Tensor, "Nr 1"]
+        weights_, trans_, _ = nerfacc.render_weight_from_density(
+            t_starts[..., 0],
+            t_ends[..., 0],
+            geo_out["density"][..., 0],
+            ray_indices=ray_indices,
+            n_rays=n_rays,
+        )
+        if self.training and self.cfg.estimator == "proposal":
+            self.vars_in_forward["trans"] = trans_.reshape(n_rays, -1)
+
+        weights = weights_[..., None]
+        opacity: Float[Tensor, "Nr 1"] = nerfacc.accumulate_along_rays(
+            weights[..., 0], values=None, ray_indices=ray_indices, n_rays=n_rays
+        )
+        depth: Float[Tensor, "Nr 1"] = nerfacc.accumulate_along_rays(
+            weights[..., 0], values=t_positions, ray_indices=ray_indices, n_rays=n_rays
+        )
+        comp_rgb_fg: Float[Tensor, "Nr Nc"] = nerfacc.accumulate_along_rays(
+            weights[..., 0], values=rgb_fg_all, ray_indices=ray_indices, n_rays=n_rays
+        )
+
+        # populate depth and opacity to each point
+        t_depth = depth[ray_indices]
+        z_variance = nerfacc.accumulate_along_rays(
+            weights[..., 0],
+            values=(t_positions - t_depth) ** 2,
+            ray_indices=ray_indices,
+            n_rays=n_rays,
+        )
+
+        if bg_color is None:
+            bg_color = comp_rgb_bg
+        else:
+            if bg_color.shape[:-1] == (batch_size,):
+                # e.g. constant random color used for Zero123
+                # [bs,3] -> [bs, 1, 1, 3]):
+                bg_color = bg_color.unsqueeze(1).unsqueeze(1)
+                #        -> [bs, height, width, 3]):
+                bg_color = bg_color.expand(-1, height, width, -1)
+
+        if bg_color.shape[:-1] == (batch_size, height, width):
+            bg_color = bg_color.reshape(batch_size * height * width, -1)
+
+        comp_rgb = comp_rgb_fg + bg_color * (1.0 - opacity)
+
+        out = {
+            "comp_rgb": comp_rgb.view(batch_size, height, width, -1),
+            "comp_rgb_fg": comp_rgb_fg.view(batch_size, height, width, -1),
+            "comp_rgb_bg": comp_rgb_bg.view(batch_size, height, width, -1),
+            "opacity": opacity.view(batch_size, height, width, 1),
+            "depth": depth.view(batch_size, height, width, 1),
+            "z_variance": z_variance.view(batch_size, height, width, 1),
+        }
+
+        if self.training:
+            out.update(
+                {
+                    "weights": weights,
+                    "t_points": t_positions,
+                    "t_intervals": t_intervals,
+                    "t_dirs": t_dirs,
+                    "ray_indices": ray_indices,
+                    "points": positions,
+                    **geo_out,
+                }
+            )
+            if "normal" in geo_out:
+                if self.cfg.return_comp_normal:
+                    comp_normal: Float[Tensor, "Nr 3"] = nerfacc.accumulate_along_rays(
+                        weights[..., 0],
+                        values=geo_out["normal"],
+                        ray_indices=ray_indices,
+                        n_rays=n_rays,
+                    )
+                    comp_normal = F.normalize(comp_normal, dim=-1)
+                    comp_normal = (
+                        (comp_normal + 1.0) / 2.0 * opacity
+                    )  # for visualization
+                    out.update(
+                        {
+                            "comp_normal": comp_normal.view(
+                                batch_size, height, width, 3
+                            ),
+                        }
+                    )
+                if self.cfg.return_normal_perturb:
+                    normal_perturb = self.geometry(
+                        positions + torch.randn_like(positions) * 1e-2,
+                        output_normal=self.material.requires_normal,
+                    )["normal"]
+                    out.update({"normal_perturb": normal_perturb})
+        else:
+            if "normal" in geo_out:
+                comp_normal = nerfacc.accumulate_along_rays(
+                    weights[..., 0],
+                    values=geo_out["normal"],
+                    ray_indices=ray_indices,
+                    n_rays=n_rays,
+                )
+                comp_normal = F.normalize(comp_normal, dim=-1)
+                comp_normal = (comp_normal + 1.0) / 2.0 * opacity  # for visualization
+                out.update(
+                    {
+                        "comp_normal": comp_normal.view(batch_size, height, width, 3),
+                    }
+                )
+
+        return out
+
+    def update_step(
+        self, epoch: int, global_step: int, on_load_weights: bool = False
+    ) -> None:
+        if self.cfg.estimator == "occgrid":
+            if self.cfg.grid_prune:
+
+                def occ_eval_fn(x):
+                    density = self.geometry.forward_density(x)
+                    # approximate for 1 - torch.exp(-density * self.render_step_size) based on taylor series
+                    return density * self.render_step_size
+
+                if self.training and not on_load_weights:
+                    self.estimator.update_every_n_steps(
+                        step=global_step, occ_eval_fn=occ_eval_fn
+                    )
+        elif self.cfg.estimator == "proposal":
+            if self.training:
+                requires_grad = self.proposal_requires_grad_fn(global_step)
+                self.vars_in_forward["requires_grad"] = requires_grad
+            else:
+                self.vars_in_forward["requires_grad"] = False
+
+    def update_step_end(self, epoch: int, global_step: int) -> None:
+        if self.cfg.estimator == "proposal" and self.training:
+            self.estimator.update_every_n_steps(
+                self.vars_in_forward["trans"],
+                self.vars_in_forward["requires_grad"],
+                loss_scaler=1.0,
+            )
+
+    def train(self, mode=True):
+        self.randomized = mode and self.cfg.randomized
+        if self.cfg.estimator == "proposal":
+            self.prop_net.train()
+        return super().train(mode=mode)
+
+    def eval(self):
+        self.randomized = False
+        if self.cfg.estimator == "proposal":
+            self.prop_net.eval()
+        return super().eval()

threestudio/models/renderers/neus_volume_renderer.py

@@ -0,0 +1,390 @@
+from dataclasses import dataclass
+from functools import partial
+
+import nerfacc
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.estimators import ImportanceEstimator
+from threestudio.models.geometry.base import BaseImplicitGeometry
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.renderers.base import VolumeRenderer
+from threestudio.utils.ops import chunk_batch, validate_empty_rays
+from threestudio.utils.typing import *
+
+
+def volsdf_density(sdf, inv_std):
+    inv_std = inv_std.clamp(0.0, 80.0)
+    beta = 1 / inv_std
+    alpha = inv_std
+    return alpha * (0.5 + 0.5 * sdf.sign() * torch.expm1(-sdf.abs() / beta))
+
+
+class LearnedVariance(nn.Module):
+    def __init__(self, init_val):
+        super(LearnedVariance, self).__init__()
+        self.register_parameter("_inv_std", nn.Parameter(torch.tensor(init_val)))
+
+    @property
+    def inv_std(self):
+        val = torch.exp(self._inv_std * 10.0)
+        return val
+
+    def forward(self, x):
+        return torch.ones_like(x) * self.inv_std.clamp(1.0e-6, 1.0e6)
+
+
+@threestudio.register("neus-volume-renderer")
+class NeuSVolumeRenderer(VolumeRenderer):
+    @dataclass
+    class Config(VolumeRenderer.Config):
+        num_samples_per_ray: int = 512
+        randomized: bool = True
+        eval_chunk_size: int = 160000
+        learned_variance_init: float = 0.3
+        cos_anneal_end_steps: int = 0
+        use_volsdf: bool = False
+
+        near_plane: float = 0.0
+        far_plane: float = 1e10
+
+        # in ['occgrid', 'importance']
+        estimator: str = "occgrid"
+
+        # for occgrid
+        grid_prune: bool = True
+        prune_alpha_threshold: bool = True
+
+        # for importance
+        num_samples_per_ray_importance: int = 64
+
+    cfg: Config
+
+    def configure(
+        self,
+        geometry: BaseImplicitGeometry,
+        material: BaseMaterial,
+        background: BaseBackground,
+    ) -> None:
+        super().configure(geometry, material, background)
+        self.variance = LearnedVariance(self.cfg.learned_variance_init)
+        if self.cfg.estimator == "occgrid":
+            self.estimator = nerfacc.OccGridEstimator(
+                roi_aabb=self.bbox.view(-1), resolution=32, levels=1
+            )
+            if not self.cfg.grid_prune:
+                self.estimator.occs.fill_(True)
+                self.estimator.binaries.fill_(True)
+            self.render_step_size = (
+                1.732 * 2 * self.cfg.radius / self.cfg.num_samples_per_ray
+            )
+            self.randomized = self.cfg.randomized
+        elif self.cfg.estimator == "importance":
+            self.estimator = ImportanceEstimator()
+        else:
+            raise NotImplementedError(
+                "unknown estimator, should be in ['occgrid', 'importance']"
+            )
+        self.cos_anneal_ratio = 1.0
+
+    def get_alpha(self, sdf, normal, dirs, dists):
+        inv_std = self.variance(sdf)
+        if self.cfg.use_volsdf:
+            alpha = torch.abs(dists.detach()) * volsdf_density(sdf, inv_std)
+        else:
+            true_cos = (dirs * normal).sum(-1, keepdim=True)
+            # "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
+            # the cos value "not dead" at the beginning training iterations, for better convergence.
+            iter_cos = -(
+                F.relu(-true_cos * 0.5 + 0.5) * (1.0 - self.cos_anneal_ratio)
+                + F.relu(-true_cos) * self.cos_anneal_ratio
+            )  # always non-positive
+
+            # Estimate signed distances at section points
+            estimated_next_sdf = sdf + iter_cos * dists * 0.5
+            estimated_prev_sdf = sdf - iter_cos * dists * 0.5
+
+            prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_std)
+            next_cdf = torch.sigmoid(estimated_next_sdf * inv_std)
+
+            p = prev_cdf - next_cdf
+            c = prev_cdf
+
+            alpha = ((p + 1e-5) / (c + 1e-5)).clip(0.0, 1.0)
+        return alpha
+
+    def forward(
+        self,
+        rays_o: Float[Tensor, "B H W 3"],
+        rays_d: Float[Tensor, "B H W 3"],
+        light_positions: Float[Tensor, "B 3"],
+        bg_color: Optional[Tensor] = None,
+        **kwargs
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        batch_size, height, width = rays_o.shape[:3]
+        rays_o_flatten: Float[Tensor, "Nr 3"] = rays_o.reshape(-1, 3)
+        rays_d_flatten: Float[Tensor, "Nr 3"] = rays_d.reshape(-1, 3)
+        light_positions_flatten: Float[Tensor, "Nr 3"] = (
+            light_positions.reshape(-1, 1, 1, 3)
+            .expand(-1, height, width, -1)
+            .reshape(-1, 3)
+        )
+        n_rays = rays_o_flatten.shape[0]
+
+        if self.cfg.estimator == "occgrid":
+
+            def alpha_fn(t_starts, t_ends, ray_indices):
+                t_starts, t_ends = t_starts[..., None], t_ends[..., None]
+                t_origins = rays_o_flatten[ray_indices]
+                t_positions = (t_starts + t_ends) / 2.0
+                t_dirs = rays_d_flatten[ray_indices]
+                positions = t_origins + t_dirs * t_positions
+                if self.training:
+                    sdf = self.geometry.forward_sdf(positions)[..., 0]
+                else:
+                    sdf = chunk_batch(
+                        self.geometry.forward_sdf,
+                        self.cfg.eval_chunk_size,
+                        positions,
+                    )[..., 0]
+
+                inv_std = self.variance(sdf)
+                if self.cfg.use_volsdf:
+                    alpha = self.render_step_size * volsdf_density(sdf, inv_std)
+                else:
+                    estimated_next_sdf = sdf - self.render_step_size * 0.5
+                    estimated_prev_sdf = sdf + self.render_step_size * 0.5
+                    prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_std)
+                    next_cdf = torch.sigmoid(estimated_next_sdf * inv_std)
+                    p = prev_cdf - next_cdf
+                    c = prev_cdf
+                    alpha = ((p + 1e-5) / (c + 1e-5)).clip(0.0, 1.0)
+
+                return alpha
+
+            if not self.cfg.grid_prune:
+                with torch.no_grad():
+                    ray_indices, t_starts_, t_ends_ = self.estimator.sampling(
+                        rays_o_flatten,
+                        rays_d_flatten,
+                        alpha_fn=None,
+                        near_plane=self.cfg.near_plane,
+                        far_plane=self.cfg.far_plane,
+                        render_step_size=self.render_step_size,
+                        alpha_thre=0.0,
+                        stratified=self.randomized,
+                        cone_angle=0.0,
+                        early_stop_eps=0,
+                    )
+            else:
+                with torch.no_grad():
+                    ray_indices, t_starts_, t_ends_ = self.estimator.sampling(
+                        rays_o_flatten,
+                        rays_d_flatten,
+                        alpha_fn=alpha_fn if self.cfg.prune_alpha_threshold else None,
+                        near_plane=self.cfg.near_plane,
+                        far_plane=self.cfg.far_plane,
+                        render_step_size=self.render_step_size,
+                        alpha_thre=0.01 if self.cfg.prune_alpha_threshold else 0.0,
+                        stratified=self.randomized,
+                        cone_angle=0.0,
+                    )
+        elif self.cfg.estimator == "importance":
+
+            def prop_sigma_fn(
+                t_starts: Float[Tensor, "Nr Ns"],
+                t_ends: Float[Tensor, "Nr Ns"],
+                proposal_network,
+            ):
+                if self.cfg.use_volsdf:
+                    t_origins: Float[Tensor, "Nr 1 3"] = rays_o_flatten.unsqueeze(-2)
+                    t_dirs: Float[Tensor, "Nr 1 3"] = rays_d_flatten.unsqueeze(-2)
+                    positions: Float[Tensor, "Nr Ns 3"] = (
+                        t_origins + t_dirs * (t_starts + t_ends)[..., None] / 2.0
+                    )
+                    with torch.no_grad():
+                        geo_out = chunk_batch(
+                            proposal_network,
+                            self.cfg.eval_chunk_size,
+                            positions.reshape(-1, 3),
+                            output_normal=False,
+                        )
+                        inv_std = self.variance(geo_out["sdf"])
+                        density = volsdf_density(geo_out["sdf"], inv_std)
+                    return density.reshape(positions.shape[:2])
+                else:
+                    raise ValueError(
+                        "Currently only VolSDF supports importance sampling."
+                    )
+
+            t_starts_, t_ends_ = self.estimator.sampling(
+                prop_sigma_fns=[partial(prop_sigma_fn, proposal_network=self.geometry)],
+                prop_samples=[self.cfg.num_samples_per_ray_importance],
+                num_samples=self.cfg.num_samples_per_ray,
+                n_rays=n_rays,
+                near_plane=self.cfg.near_plane,
+                far_plane=self.cfg.far_plane,
+                sampling_type="uniform",
+                stratified=self.randomized,
+            )
+            ray_indices = (
+                torch.arange(n_rays, device=rays_o_flatten.device)
+                .unsqueeze(-1)
+                .expand(-1, t_starts_.shape[1])
+            )
+            ray_indices = ray_indices.flatten()
+            t_starts_ = t_starts_.flatten()
+            t_ends_ = t_ends_.flatten()
+        else:
+            raise NotImplementedError
+
+        ray_indices, t_starts_, t_ends_ = validate_empty_rays(
+            ray_indices, t_starts_, t_ends_
+        )
+        ray_indices = ray_indices.long()
+        t_starts, t_ends = t_starts_[..., None], t_ends_[..., None]
+        t_origins = rays_o_flatten[ray_indices]
+        t_dirs = rays_d_flatten[ray_indices]
+        t_light_positions = light_positions_flatten[ray_indices]
+        t_positions = (t_starts + t_ends) / 2.0
+        positions = t_origins + t_dirs * t_positions
+        t_intervals = t_ends - t_starts
+
+        if self.training:
+            geo_out = self.geometry(positions, output_normal=True)
+            rgb_fg_all = self.material(
+                viewdirs=t_dirs,
+                positions=positions,
+                light_positions=t_light_positions,
+                **geo_out,
+                **kwargs
+            )
+            comp_rgb_bg = self.background(dirs=rays_d)
+        else:
+            geo_out = chunk_batch(
+                self.geometry,
+                self.cfg.eval_chunk_size,
+                positions,
+                output_normal=True,
+            )
+            rgb_fg_all = chunk_batch(
+                self.material,
+                self.cfg.eval_chunk_size,
+                viewdirs=t_dirs,
+                positions=positions,
+                light_positions=t_light_positions,
+                **geo_out
+            )
+            comp_rgb_bg = chunk_batch(
+                self.background, self.cfg.eval_chunk_size, dirs=rays_d
+            )
+
+        # grad or normal?
+        alpha: Float[Tensor, "Nr 1"] = self.get_alpha(
+            geo_out["sdf"], geo_out["normal"], t_dirs, t_intervals
+        )
+
+        weights: Float[Tensor, "Nr 1"]
+        weights_, _ = nerfacc.render_weight_from_alpha(
+            alpha[..., 0],
+            ray_indices=ray_indices,
+            n_rays=n_rays,
+        )
+        weights = weights_[..., None]
+        opacity: Float[Tensor, "Nr 1"] = nerfacc.accumulate_along_rays(
+            weights[..., 0], values=None, ray_indices=ray_indices, n_rays=n_rays
+        )
+        depth: Float[Tensor, "Nr 1"] = nerfacc.accumulate_along_rays(
+            weights[..., 0], values=t_positions, ray_indices=ray_indices, n_rays=n_rays
+        )
+        comp_rgb_fg: Float[Tensor, "Nr Nc"] = nerfacc.accumulate_along_rays(
+            weights[..., 0], values=rgb_fg_all, ray_indices=ray_indices, n_rays=n_rays
+        )
+
+        if bg_color is None:
+            bg_color = comp_rgb_bg
+
+        if bg_color.shape[:-1] == (batch_size, height, width):
+            bg_color = bg_color.reshape(batch_size * height * width, -1)
+
+        comp_rgb = comp_rgb_fg + bg_color * (1.0 - opacity)
+
+        out = {
+            "comp_rgb": comp_rgb.view(batch_size, height, width, -1),
+            "comp_rgb_fg": comp_rgb_fg.view(batch_size, height, width, -1),
+            "comp_rgb_bg": comp_rgb_bg.view(batch_size, height, width, -1),
+            "opacity": opacity.view(batch_size, height, width, 1),
+            "depth": depth.view(batch_size, height, width, 1),
+        }
+
+        if self.training:
+            out.update(
+                {
+                    "weights": weights,
+                    "t_points": t_positions,
+                    "t_intervals": t_intervals,
+                    "t_dirs": t_dirs,
+                    "ray_indices": ray_indices,
+                    "points": positions,
+                    **geo_out,
+                }
+            )
+        else:
+            if "normal" in geo_out:
+                comp_normal: Float[Tensor, "Nr 3"] = nerfacc.accumulate_along_rays(
+                    weights[..., 0],
+                    values=geo_out["normal"],
+                    ray_indices=ray_indices,
+                    n_rays=n_rays,
+                )
+                comp_normal = F.normalize(comp_normal, dim=-1)
+                comp_normal = (comp_normal + 1.0) / 2.0 * opacity  # for visualization
+                out.update(
+                    {
+                        "comp_normal": comp_normal.view(batch_size, height, width, 3),
+                    }
+                )
+        out.update({"inv_std": self.variance.inv_std})
+        return out
+
+    def update_step(
+        self, epoch: int, global_step: int, on_load_weights: bool = False
+    ) -> None:
+        self.cos_anneal_ratio = (
+            1.0
+            if self.cfg.cos_anneal_end_steps == 0
+            else min(1.0, global_step / self.cfg.cos_anneal_end_steps)
+        )
+        if self.cfg.estimator == "occgrid":
+            if self.cfg.grid_prune:
+
+                def occ_eval_fn(x):
+                    sdf = self.geometry.forward_sdf(x)
+                    inv_std = self.variance(sdf)
+                    if self.cfg.use_volsdf:
+                        alpha = self.render_step_size * volsdf_density(sdf, inv_std)
+                    else:
+                        estimated_next_sdf = sdf - self.render_step_size * 0.5
+                        estimated_prev_sdf = sdf + self.render_step_size * 0.5
+                        prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_std)
+                        next_cdf = torch.sigmoid(estimated_next_sdf * inv_std)
+                        p = prev_cdf - next_cdf
+                        c = prev_cdf
+                        alpha = ((p + 1e-5) / (c + 1e-5)).clip(0.0, 1.0)
+                    return alpha
+
+                if self.training and not on_load_weights:
+                    self.estimator.update_every_n_steps(
+                        step=global_step, occ_eval_fn=occ_eval_fn
+                    )
+
+    def train(self, mode=True):
+        self.randomized = mode and self.cfg.randomized
+        return super().train(mode=mode)
+
+    def eval(self):
+        self.randomized = False
+        return super().eval()

threestudio/models/renderers/nvdiff_rasterizer.py

@@ -0,0 +1,188 @@
+from dataclasses import dataclass
+
+import nerfacc
+import torch
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.geometry.base import BaseImplicitGeometry
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.renderers.base import Rasterizer, VolumeRenderer
+from threestudio.utils.misc import get_device
+from threestudio.utils.rasterize import NVDiffRasterizerContext
+from threestudio.utils.typing import *
+
+
+@threestudio.register("nvdiff-rasterizer")
+class NVDiffRasterizer(Rasterizer):
+    @dataclass
+    class Config(VolumeRenderer.Config):
+        context_type: str = "gl"
+
+    cfg: Config
+
+    def configure(
+        self,
+        geometry: BaseImplicitGeometry,
+        material: BaseMaterial,
+        background: BaseBackground,
+    ) -> None:
+        super().configure(geometry, material, background)
+        self.ctx = NVDiffRasterizerContext(self.cfg.context_type, get_device())
+
+    def forward(
+        self,
+        mvp_mtx: Float[Tensor, "B 4 4"],
+        camera_positions: Float[Tensor, "B 3"],
+        light_positions: Float[Tensor, "B 3"],
+        height: int,
+        width: int,
+        render_rgb: bool = True,
+        render_mask: bool = False,
+        **kwargs
+    ) -> Dict[str, Any]:
+        batch_size = mvp_mtx.shape[0]
+        mesh = self.geometry.isosurface()
+
+        v_pos_clip: Float[Tensor, "B Nv 4"] = self.ctx.vertex_transform(
+            mesh.v_pos, mvp_mtx
+        )
+        rast, _ = self.ctx.rasterize(v_pos_clip, mesh.t_pos_idx, (height, width))
+        mask = rast[..., 3:] > 0
+        mask_aa = self.ctx.antialias(mask.float(), rast, v_pos_clip, mesh.t_pos_idx)
+
+        out = {"opacity": mask_aa, "mesh": mesh}
+
+        if render_mask:
+            # get front-view visibility mask
+            with torch.no_grad():
+                mvp_mtx_ref = kwargs["mvp_mtx_ref"] # FIXME
+                v_pos_clip_front: Float[Tensor, "B Nv 4"] = self.ctx.vertex_transform(
+                    mesh.v_pos, mvp_mtx_ref
+                )
+                rast_front, _ = self.ctx.rasterize(v_pos_clip_front, mesh.t_pos_idx, (height, width))
+                mask_front = rast_front[..., 3:]
+                mask_front = mask_front[mask_front > 0] - 1.
+                faces_vis = mesh.t_pos_idx[mask_front.long()]
+
+                mesh._v_rgb = torch.zeros(mesh.v_pos.shape[0], 1).to(mesh.v_pos)
+                mesh._v_rgb[faces_vis[:,0]] = 1.
+                mesh._v_rgb[faces_vis[:,1]] = 1.
+                mesh._v_rgb[faces_vis[:,2]] = 1.
+                mask_vis, _ = self.ctx.interpolate_one(mesh._v_rgb, rast, mesh.t_pos_idx)
+                mask_vis = mask_vis > 0.
+                # from torchvision.utils import save_image
+                # save_image(mask_vis.permute(0,3,1,2).float(), "debug.png")
+                out.update({"mask": 1.0 - mask_vis.float()})
+                
+                # FIXME: paste texture back to mesh
+                # import cv2
+                # import imageio
+                # import numpy as np
+
+                # gt_rgb = imageio.imread("load/images/tiger_nurse_rgba.png")/255.
+                # gt_rgb = cv2.resize(gt_rgb[:,:,:3],(512, 512))
+                # gt_rgb = torch.Tensor(gt_rgb[None,...]).permute(0,3,1,2).to(v_pos_clip_front)
+
+                # # align to up-z and front-x
+                # dir2vec = {
+                #     "+x": np.array([1, 0, 0]),
+                #     "+y": np.array([0, 1, 0]),
+                #     "+z": np.array([0, 0, 1]),
+                #     "-x": np.array([-1, 0, 0]),
+                #     "-y": np.array([0, -1, 0]),
+                #     "-z": np.array([0, 0, -1]),
+                # }
+                # z_, x_ = (
+                #     dir2vec["-y"],
+                #     dir2vec["-z"],
+                # )
+
+                # y_ = np.cross(z_, x_)
+                # std2mesh = np.stack([x_, y_, z_], axis=0).T
+                # v_pos_ = (torch.mm(torch.tensor(std2mesh).to(mesh.v_pos), mesh.v_pos.T).T) * 2
+                # print(v_pos_.min(), v_pos_.max())
+
+                # mesh._v_rgb=F.grid_sample(gt_rgb, v_pos_[None, None][..., :2], mode="nearest").permute(3,1,0,2).squeeze(-1).squeeze(-1).contiguous()
+                # rgb_vis, _ = self.ctx.interpolate_one(mesh._v_rgb, rast, mesh.t_pos_idx)
+                # rgb_vis_aa = self.ctx.antialias(
+                #     rgb_vis, rast, v_pos_clip, mesh.t_pos_idx
+                # )
+                # from torchvision.utils import save_image
+                # save_image(rgb_vis_aa.permute(0,3,1,2), "debug.png")
+        
+
+        gb_normal, _ = self.ctx.interpolate_one(mesh.v_nrm, rast, mesh.t_pos_idx)
+        gb_normal = F.normalize(gb_normal, dim=-1)
+        gb_normal_aa = torch.lerp(
+            torch.zeros_like(gb_normal), (gb_normal + 1.0) / 2.0, mask.float()
+        )
+        gb_normal_aa = self.ctx.antialias(
+            gb_normal_aa, rast, v_pos_clip, mesh.t_pos_idx
+        )
+        out.update({"comp_normal": gb_normal_aa})  # in [0, 1]
+
+        # Compute normal in view space.
+        # TODO: make is clear whether to compute this.
+        w2c = kwargs["c2w"][:, :3, :3].inverse()
+        gb_normal_viewspace = torch.einsum("bij,bhwj->bhwi", w2c, gb_normal)
+        gb_normal_viewspace = F.normalize(gb_normal_viewspace, dim=-1)
+        bg_normal = torch.zeros_like(gb_normal_viewspace)
+        bg_normal[..., 2] = 1
+        gb_normal_viewspace_aa = torch.lerp(
+            (bg_normal + 1.0) / 2.0,
+            (gb_normal_viewspace + 1.0) / 2.0,
+            mask.float(),
+        ).contiguous()
+        gb_normal_viewspace_aa = self.ctx.antialias(
+            gb_normal_viewspace_aa, rast, v_pos_clip, mesh.t_pos_idx
+        )
+        out.update({"comp_normal_viewspace": gb_normal_viewspace_aa})
+
+        # TODO: make it clear whether to compute the normal, now we compute it in all cases
+        # consider using: require_normal_computation = render_normal or (render_rgb and material.requires_normal)
+        # or
+        # render_normal = render_normal or (render_rgb and material.requires_normal)
+
+        if render_rgb:
+            selector = mask[..., 0]
+
+            gb_pos, _ = self.ctx.interpolate_one(mesh.v_pos, rast, mesh.t_pos_idx)
+            gb_viewdirs = F.normalize(
+                gb_pos - camera_positions[:, None, None, :], dim=-1
+            )
+            gb_light_positions = light_positions[:, None, None, :].expand(
+                -1, height, width, -1
+            )
+
+            positions = gb_pos[selector]
+            geo_out = self.geometry(positions, output_normal=False)
+
+            extra_geo_info = {}
+            if self.material.requires_normal:
+                extra_geo_info["shading_normal"] = gb_normal[selector]
+            if self.material.requires_tangent:
+                gb_tangent, _ = self.ctx.interpolate_one(
+                    mesh.v_tng, rast, mesh.t_pos_idx
+                )
+                gb_tangent = F.normalize(gb_tangent, dim=-1)
+                extra_geo_info["tangent"] = gb_tangent[selector]
+
+            rgb_fg = self.material(
+                viewdirs=gb_viewdirs[selector],
+                positions=positions,
+                light_positions=gb_light_positions[selector],
+                **extra_geo_info,
+                **geo_out
+            )
+            gb_rgb_fg = torch.zeros(batch_size, height, width, 3).to(rgb_fg)
+            gb_rgb_fg[selector] = rgb_fg
+
+            gb_rgb_bg = self.background(dirs=gb_viewdirs)
+            gb_rgb = torch.lerp(gb_rgb_bg, gb_rgb_fg, mask.float())
+            gb_rgb_aa = self.ctx.antialias(gb_rgb, rast, v_pos_clip, mesh.t_pos_idx)
+
+            out.update({"comp_rgb": gb_rgb_aa, "comp_rgb_bg": gb_rgb_bg})
+
+        return out

threestudio/models/renderers/patch_renderer.py

@@ -0,0 +1,106 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.geometry.base import BaseImplicitGeometry
+from threestudio.models.materials.base import BaseMaterial
+from threestudio.models.renderers.base import VolumeRenderer
+from threestudio.utils.typing import *
+
+
+@threestudio.register("patch-renderer")
+class PatchRenderer(VolumeRenderer):
+    @dataclass
+    class Config(VolumeRenderer.Config):
+        patch_size: int = 128
+        base_renderer_type: str = ""
+        base_renderer: Optional[VolumeRenderer.Config] = None
+        global_detach: bool = False
+        global_downsample: int = 4
+
+    cfg: Config
+
+    def configure(
+        self,
+        geometry: BaseImplicitGeometry,
+        material: BaseMaterial,
+        background: BaseBackground,
+    ) -> None:
+        self.base_renderer = threestudio.find(self.cfg.base_renderer_type)(
+            self.cfg.base_renderer,
+            geometry=geometry,
+            material=material,
+            background=background,
+        )
+
+    def forward(
+        self,
+        rays_o: Float[Tensor, "B H W 3"],
+        rays_d: Float[Tensor, "B H W 3"],
+        light_positions: Float[Tensor, "B 3"],
+        bg_color: Optional[Tensor] = None,
+        **kwargs
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        B, H, W, _ = rays_o.shape
+
+        if self.base_renderer.training:
+            downsample = self.cfg.global_downsample
+            global_rays_o = torch.nn.functional.interpolate(
+                rays_o.permute(0, 3, 1, 2),
+                (H // downsample, W // downsample),
+                mode="bilinear",
+            ).permute(0, 2, 3, 1)
+            global_rays_d = torch.nn.functional.interpolate(
+                rays_d.permute(0, 3, 1, 2),
+                (H // downsample, W // downsample),
+                mode="bilinear",
+            ).permute(0, 2, 3, 1)
+            out_global = self.base_renderer(
+                global_rays_o, global_rays_d, light_positions, bg_color, **kwargs
+            )
+
+            PS = self.cfg.patch_size
+            patch_x = torch.randint(0, W - PS, (1,)).item()
+            patch_y = torch.randint(0, H - PS, (1,)).item()
+            patch_rays_o = rays_o[:, patch_y : patch_y + PS, patch_x : patch_x + PS]
+            patch_rays_d = rays_d[:, patch_y : patch_y + PS, patch_x : patch_x + PS]
+            out = self.base_renderer(
+                patch_rays_o, patch_rays_d, light_positions, bg_color, **kwargs
+            )
+
+            valid_patch_key = []
+            for key in out:
+                if torch.is_tensor(out[key]):
+                    if len(out[key].shape) == len(out["comp_rgb"].shape):
+                        if out[key][..., 0].shape == out["comp_rgb"][..., 0].shape:
+                            valid_patch_key.append(key)
+            for key in valid_patch_key:
+                out_global[key] = F.interpolate(
+                    out_global[key].permute(0, 3, 1, 2), (H, W), mode="bilinear"
+                ).permute(0, 2, 3, 1)
+                if self.cfg.global_detach:
+                    out_global[key] = out_global[key].detach()
+                out_global[key][
+                    :, patch_y : patch_y + PS, patch_x : patch_x + PS
+                ] = out[key]
+            out = out_global
+        else:
+            out = self.base_renderer(
+                rays_o, rays_d, light_positions, bg_color, **kwargs
+            )
+
+        return out
+
+    def update_step(
+        self, epoch: int, global_step: int, on_load_weights: bool = False
+    ) -> None:
+        self.base_renderer.update_step(epoch, global_step, on_load_weights)
+
+    def train(self, mode=True):
+        return self.base_renderer.train(mode)
+
+    def eval(self):
+        return self.base_renderer.eval()