chores: rebase commits

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