chores: rebase commits

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()