Delete lrm

lrm/__init__.py

@@ -1,5 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.

lrm/cam_utils.py

@@ -1,138 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-
-import torch
-import numpy as np
-import math
-
-"""
-R: (N, 3, 3)
-T: (N, 3)
-E: (N, 4, 4)
-vector: (N, 3)
-"""
-
-
-def compose_extrinsic_R_T(R: torch.Tensor, T: torch.Tensor):
-    """
-    Compose the standard form extrinsic matrix from R and T.
-    Batched I/O.
-    """
-    RT = torch.cat((R, T.unsqueeze(-1)), dim=-1)
-    return compose_extrinsic_RT(RT)
-
-
-def compose_extrinsic_RT(RT: torch.Tensor):
-    """
-    Compose the standard form extrinsic matrix from RT.
-    Batched I/O.
-    """
-    return torch.cat([
-        RT,
-        torch.tensor([[[0, 0, 0, 1]]], dtype=torch.float32).repeat(RT.shape[0], 1, 1).to(RT.device)
-        ], dim=1)
-
-
-def decompose_extrinsic_R_T(E: torch.Tensor):
-    """
-    Decompose the standard extrinsic matrix into R and T.
-    Batched I/O.
-    """
-    RT = decompose_extrinsic_RT(E)
-    return RT[:, :, :3], RT[:, :, 3]
-
-
-def decompose_extrinsic_RT(E: torch.Tensor):
-    """
-    Decompose the standard extrinsic matrix into RT.
-    Batched I/O.
-    """
-    return E[:, :3, :]
-
-
-def get_normalized_camera_intrinsics(intrinsics: torch.Tensor):
-    """
-    intrinsics: (N, 3, 2), [[fx, fy], [cx, cy], [width, height]]
-    Return batched fx, fy, cx, cy
-    """
-    fx, fy = intrinsics[:, 0, 0], intrinsics[:, 0, 1]
-    cx, cy = intrinsics[:, 1, 0], intrinsics[:, 1, 1]
-    width, height = intrinsics[:, 2, 0], intrinsics[:, 2, 1]
-    fx, fy = fx / width, fy / height
-    cx, cy = cx / width, cy / height
-    return fx, fy, cx, cy
-
-
-def build_camera_principle(RT: torch.Tensor, intrinsics: torch.Tensor):
-    """
-    RT: (N, 3, 4)
-    intrinsics: (N, 3, 2), [[fx, fy], [cx, cy], [width, height]]
-    """
-    fx, fy, cx, cy = get_normalized_camera_intrinsics(intrinsics)
-    return torch.cat([
-        RT.reshape(-1, 12),
-        fx.unsqueeze(-1), fy.unsqueeze(-1), cx.unsqueeze(-1), cy.unsqueeze(-1),
-    ], dim=-1)
-
-
-def build_camera_standard(RT: torch.Tensor, intrinsics: torch.Tensor):
-    """
-    RT: (N, 3, 4)
-    intrinsics: (N, 3, 2), [[fx, fy], [cx, cy], [width, height]]
-    """
-    E = compose_extrinsic_RT(RT)
-    fx, fy, cx, cy = get_normalized_camera_intrinsics(intrinsics)
-    I = torch.stack([
-        torch.stack([fx, torch.zeros_like(fx), cx], dim=-1),
-        torch.stack([torch.zeros_like(fy), fy, cy], dim=-1),
-        torch.tensor([[0, 0, 1]], dtype=torch.float32, device=RT.device).repeat(RT.shape[0], 1),
-    ], dim=1)
-    return torch.cat([
-        E.reshape(-1, 16),
-        I.reshape(-1, 9),
-    ], dim=-1)
-
-
-def center_looking_at_camera_pose(camera_position: torch.Tensor, look_at: torch.Tensor = None, up_world: torch.Tensor = None):
-    """
-    camera_position: (M, 3)
-    look_at: (3)
-    up_world: (3)
-    return: (M, 3, 4)
-    """
-    # by default, looking at the origin and world up is pos-z
-    if look_at is None:
-        look_at = torch.tensor([0, 0, 0], dtype=torch.float32)
-    if up_world is None:
-        up_world = torch.tensor([0, 0, 1], dtype=torch.float32)
-    look_at = look_at.unsqueeze(0).repeat(camera_position.shape[0], 1)
-    up_world = up_world.unsqueeze(0).repeat(camera_position.shape[0], 1)
-
-    z_axis = camera_position - look_at
-    z_axis = z_axis / z_axis.norm(dim=-1, keepdim=True)
-    x_axis = torch.cross(up_world, z_axis)
-    x_axis = x_axis / x_axis.norm(dim=-1, keepdim=True)
-    y_axis = torch.cross(z_axis, x_axis)
-    y_axis = y_axis / y_axis.norm(dim=-1, keepdim=True)
-    extrinsics = torch.stack([x_axis, y_axis, z_axis, camera_position], dim=-1)
-    return extrinsics
-
-def get_surrounding_views(M, radius, elevation):
-#   convert spherical coordinates (radius, azimuth, elevation) to Cartesian coordinates (x, y, z).
-    camera_positions = []
-    rand_theta= np.random.uniform(0, np.pi/180)
-    elevation = math.radians(elevation)
-    for i in range(M):
-        theta = 2 * math.pi * i / M  + rand_theta
-        x = radius * math.cos(theta) * math.cos(elevation)
-        y = radius * math.sin(theta) * math.cos(elevation)
-        z =  radius * math.sin(elevation)
-        camera_positions.append([x, y, z])
-    camera_positions = torch.tensor(camera_positions, dtype=torch.float32)
-    extrinsics = center_looking_at_camera_pose(camera_positions)
-
-    return extrinsics

lrm/inferrer.py

@@ -1,232 +0,0 @@
-import torch
-import math
-import os
-import imageio
-import mcubes
-import trimesh
-import numpy as np
-import argparse
-from torchvision.utils import save_image
-from PIL import Image
-import glob
-from .models.generator import LRMGenerator  # Make sure this import is correct
-from .cam_utils import build_camera_principle, build_camera_standard, center_looking_at_camera_pose  # Make sure this import is correct
-from functools import partial
-from rembg import remove, new_session
-from kiui.op import recenter
-import kiui
-
-class LRMInferrer:
-    def __init__(self, model_name: str, resume: str):
-        print("Initializing LRMInferrer")
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        _model_kwargs = {'camera_embed_dim': 1024, 'rendering_samples_per_ray': 128, 'transformer_dim': 1024, 'transformer_layers': 16, 'transformer_heads': 16, 'triplane_low_res': 32, 'triplane_high_res': 64, 'triplane_dim': 80, 'encoder_freeze': False}
-        
-        self.model = self._build_model(_model_kwargs).eval().to(self.device)
-        checkpoint = torch.load(resume, map_location='cpu')
-        state_dict = checkpoint['model_state_dict']
-        self.model.load_state_dict(state_dict)
-        del checkpoint, state_dict
-        torch.cuda.empty_cache()
-
-    def __enter__(self):
-        print("Entering context")
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        print("Exiting context")
-        if exc_type:
-            print(f"Exception type: {exc_type}")
-            print(f"Exception value: {exc_val}")
-            print(f"Traceback: {exc_tb}")
-
-    def _build_model(self, model_kwargs):
-        print("Building model")
-        model = LRMGenerator(**model_kwargs).to(self.device)
-        print("Loaded model from checkpoint")
-        return model
-
-    @staticmethod
-    def get_surrounding_views(M, radius, elevation):
-        camera_positions = []
-        rand_theta = np.random.uniform(0, np.pi/180)
-        elevation = math.radians(elevation)
-        for i in range(M):
-            theta = 2 * math.pi * i / M + rand_theta
-            x = radius * math.cos(theta) * math.cos(elevation)
-            y = radius * math.sin(theta) * math.cos(elevation)
-            z = radius * math.sin(elevation)
-            camera_positions.append([x, y, z])
-        camera_positions = torch.tensor(camera_positions, dtype=torch.float32)
-        extrinsics = center_looking_at_camera_pose(camera_positions)
-        return extrinsics
-
-    @staticmethod
-    def _default_intrinsics():
-        fx = fy = 384
-        cx = cy = 256
-        w = h = 512
-        intrinsics = torch.tensor([
-            [fx, fy],
-            [cx, cy],
-            [w, h],
-        ], dtype=torch.float32)
-        return intrinsics
-
-    def _default_source_camera(self, batch_size: int = 1):
-        dist_to_center = 1.5
-        canonical_camera_extrinsics = torch.tensor([[
-            [0, 0, 1, 1],
-            [1, 0, 0, 0],
-            [0, 1, 0, 0],
-        ]], dtype=torch.float32)
-        canonical_camera_intrinsics = self._default_intrinsics().unsqueeze(0)
-        source_camera = build_camera_principle(canonical_camera_extrinsics, canonical_camera_intrinsics)
-        return source_camera.repeat(batch_size, 1)
-
-    def _default_render_cameras(self, batch_size: int = 1):
-        render_camera_extrinsics = self.get_surrounding_views(160, 1.5, 0)
-        render_camera_intrinsics = self._default_intrinsics().unsqueeze(0).repeat(render_camera_extrinsics.shape[0], 1, 1)
-        render_cameras = build_camera_standard(render_camera_extrinsics, render_camera_intrinsics)
-        return render_cameras.unsqueeze(0).repeat(batch_size, 1, 1)
-
-    @staticmethod
-    def images_to_video(images, output_path, fps, verbose=False):
-        os.makedirs(os.path.dirname(output_path), exist_ok=True)
-        frames = []
-        for i in range(images.shape[0]):
-            frame = (images[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
-            assert frame.shape[0] == images.shape[2] and frame.shape[1] == images.shape[3], \
-                f"Frame shape mismatch: {frame.shape} vs {images.shape}"
-            assert frame.min() >= 0 and frame.max() <= 255, \
-                f"Frame value out of range: {frame.min()} ~ {frame.max()}"
-            frames.append(frame)
-        imageio.mimwrite(output_path, np.stack(frames), fps=fps)
-        if verbose:
-            print(f"Saved video to {output_path}")
-
-    def infer_single(self, image: torch.Tensor, render_size: int, mesh_size: int, export_video: bool, export_mesh: bool):
-        print("infer_single called")
-        mesh_thres = 1.0
-        chunk_size = 2
-        batch_size = 1
-
-        source_camera = self._default_source_camera(batch_size).to(self.device)
-        render_cameras = self._default_render_cameras(batch_size).to(self.device)
-
-        with torch.no_grad():
-            planes = self.model.forward(image, source_camera)
-            results = {}
-
-            if export_video:
-                print("Starting export_video")
-                frames = []
-                for i in range(0, render_cameras.shape[1], chunk_size):
-                    print(f"Processing chunk {i} to {i + chunk_size}")
-                    frames.append(
-                        self.model.synthesizer(
-                            planes,
-                            render_cameras[:, i:i+chunk_size],
-                            render_size,
-                            render_size,
-                            0,
-                            0
-                        )
-                    )
-                frames = {
-                    k: torch.cat([r[k] for r in frames], dim=1)
-                    for k in frames[0].keys()
-                }
-                results.update({
-                    'frames': frames,
-                })
-                print("Finished export_video")
-
-            if export_mesh:
-                print("Starting export_mesh")
-                grid_out = self.model.synthesizer.forward_grid(
-                    planes=planes,
-                    grid_size=mesh_size,
-                )
-                vtx, faces = mcubes.marching_cubes(grid_out['sigma'].float().squeeze(0).squeeze(-1).cpu().numpy(), mesh_thres)
-                vtx = vtx / (mesh_size - 1) * 2 - 1
-                vtx_tensor = torch.tensor(vtx, dtype=torch.float32, device=self.device).unsqueeze(0)
-                vtx_colors = self.model.synthesizer.forward_points(planes, vtx_tensor)['rgb'].float().squeeze(0).cpu().numpy()
-                vtx_colors = (vtx_colors * 255).astype(np.uint8)
-                mesh = trimesh.Trimesh(vertices=vtx, faces=faces, vertex_colors=vtx_colors)
-                results.update({
-                    'mesh': mesh,
-                })
-                print("Finished export_mesh")
-
-            return results
-
-    def infer(self, source_image: str, dump_path: str, source_size: int, render_size: int, mesh_size: int, export_video: bool, export_mesh: bool):
-        print("infer called")
-        session = new_session("isnet-general-use")
-        rembg_remove = partial(remove, session=session)
-        image_name = os.path.basename(source_image)
-        uid = image_name.split('.')[0]
-
-        image = kiui.read_image(source_image, mode='uint8')
-        image = rembg_remove(image)
-        mask = rembg_remove(image, only_mask=True)
-        image = recenter(image, mask, border_ratio=0.20)
-        os.makedirs(dump_path, exist_ok=True)
-
-        image = torch.tensor(np.array(image)).permute(2, 0, 1).unsqueeze(0) / 255.0
-        if image.shape[1] == 4:
-            image = image[:, :3, ...] * image[:, 3:, ...] + (1 - image[:, 3:, ...])
-        image = torch.nn.functional.interpolate(image, size=(source_size, source_size), mode='bicubic', align_corners=True)
-        image = torch.clamp(image, 0, 1)
-        save_image(image, os.path.join(dump_path, f'{uid}.png'))
-
-        results = self.infer_single(
-            image.cuda(),
-            render_size=render_size,
-            mesh_size=mesh_size,
-            export_video=export_video,
-            export_mesh=export_mesh,
-        )
-
-        if 'frames' in results:
-            renderings = results['frames']
-            for k, v in renderings.items():
-                if k == 'images_rgb':
-                    self.images_to_video(
-                        v[0],
-                        os.path.join(dump_path, f'{uid}.mp4'),
-                        fps=40,
-                    )
-                    print(f"Export video success to {dump_path}")
-
-        if 'mesh' in results:
-            mesh = results['mesh']
-            mesh.export(os.path.join(dump_path, f'{uid}.obj'), 'obj')
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--model_name', type=str, default='lrm-base-obj-v1')
-    parser.add_argument('--source_path', type=str, default='./assets/cat.png')
-    parser.add_argument('--dump_path', type=str, default='./results/single_image')
-    parser.add_argument('--source_size', type=int, default=512)
-    parser.add_argument('--render_size', type=int, default=384)
-    parser.add_argument('--mesh_size', type=int, default=512)
-    parser.add_argument('--export_video', action='store_true')
-    parser.add_argument('--export_mesh', action='store_true')
-    parser.add_argument('--resume', type=str, required=True, help='Path to a checkpoint to resume training from')
-    args = parser.parse_args()
-
-    with LRMInferrer(model_name=args.model_name, resume=args.resume) as inferrer:
-        with torch.autocast(device_type="cuda", cache_enabled=False, dtype=torch.float32):
-            print("Start inference for image:", args.source_path)
-            inferrer.infer(
-                source_image=args.source_path,
-                dump_path=args.dump_path,
-                source_size=args.source_size,
-                render_size=args.render_size,
-                mesh_size=args.mesh_size,
-                export_video=args.export_video,
-                export_mesh=args.export_mesh,
-            )
-            print("Finished inference for image:", args.source_path)

lrm/models/__init__.py

@@ -1,5 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.

lrm/models/encoders/__init__.py

@@ -1,5 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.

lrm/models/encoders/dino_wrapper2.py

@@ -1,51 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-
-import torch.nn as nn
-from transformers import ViTImageProcessor, ViTModel, AutoImageProcessor, AutoModel, Dinov2Model
-
-class DinoWrapper(nn.Module):
-    """
-    Dino v1 wrapper using huggingface transformer implementation.
-    """
-    def __init__(self, model_name: str, freeze: bool = True):
-        super().__init__()
-        self.model, self.processor = self._build_dino(model_name)
-        if freeze:
-            self._freeze()
-
-    def forward(self, image):
-        # image: [N, C, H, W], on cpu
-        # RGB image with [0,1] scale and properly sized
-        inputs = self.processor(images=image.float(), return_tensors="pt", do_rescale=False, do_resize=False).to(self.model.device)
-        # This resampling of positional embedding uses bicubic interpolation
-        outputs = self.model(**inputs)
-        last_hidden_states = outputs.last_hidden_state
-        return last_hidden_states
-
-    def _freeze(self):
-        print(f"======== Freezing DinoWrapper ========")
-        self.model.eval()
-        for name, param in self.model.named_parameters():
-            param.requires_grad = False
-
-    @staticmethod
-    def _build_dino(model_name: str, proxy_error_retries: int = 3, proxy_error_cooldown: int = 5):
-        import requests
-        try:
-            processor = AutoImageProcessor.from_pretrained('facebook/dinov2-base')
-            processor.do_center_crop = False
-            model = AutoModel.from_pretrained('facebook/dinov2-base')
-            return model, processor
-        except requests.exceptions.ProxyError as err:
-            if proxy_error_retries > 0:
-                print(f"Huggingface ProxyError: Retrying in {proxy_error_cooldown} seconds...")
-                import time
-                time.sleep(proxy_error_cooldown)
-                return DinoWrapper._build_dino(model_name, proxy_error_retries - 1, proxy_error_cooldown)
-            else:
-                raise err

lrm/models/generator.py

@@ -1,87 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-
-import torch.nn as nn
-
-from .encoders.dino_wrapper2 import DinoWrapper
-from .transformer import TriplaneTransformer
-from .rendering.synthesizer_part import TriplaneSynthesizer
-
-
-class CameraEmbedder(nn.Module):
-    """
-    Embed camera features to a high-dimensional vector.
-    
-    Reference:
-    DiT: https://github.com/facebookresearch/DiT/blob/main/models.py#L27
-    """
-    def __init__(self, raw_dim: int, embed_dim: int):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(raw_dim, embed_dim),
-            nn.SiLU(),
-            nn.Linear(embed_dim, embed_dim),
-        )
-
-    def forward(self, x):
-        return self.mlp(x)
-
-
-class LRMGenerator(nn.Module):
-    """
-    Full model of the large reconstruction model.
-    """
-    def __init__(self, camera_embed_dim: int, rendering_samples_per_ray: int,
-                 transformer_dim: int, transformer_layers: int, transformer_heads: int,
-                 triplane_low_res: int, triplane_high_res: int, triplane_dim: int,
-                 encoder_freeze: bool = True, encoder_model_name: str = 'facebook/dinov2-base', encoder_feat_dim: int = 768):
-        super().__init__()
-        
-        # attributes
-        self.encoder_feat_dim = encoder_feat_dim
-        self.camera_embed_dim = camera_embed_dim
-
-        # modules
-        self.encoder = DinoWrapper(
-            model_name=encoder_model_name,
-            freeze=encoder_freeze,
-        )
-        self.camera_embedder = CameraEmbedder(
-            raw_dim=12+4, embed_dim=camera_embed_dim,
-        )
-        self.transformer = TriplaneTransformer(
-            inner_dim=transformer_dim, num_layers=transformer_layers, num_heads=transformer_heads,
-            image_feat_dim=encoder_feat_dim,
-            camera_embed_dim=camera_embed_dim,
-            triplane_low_res=triplane_low_res, triplane_high_res=triplane_high_res, triplane_dim=triplane_dim,
-        )
-        self.synthesizer = TriplaneSynthesizer(
-            triplane_dim=triplane_dim, samples_per_ray=rendering_samples_per_ray,
-        )
-
-    def forward(self, image, camera):
-        # image: [N, C_img, H_img, W_img]
-        # camera: [N, D_cam_raw]
-        assert image.shape[0] == camera.shape[0], "Batch size mismatch for image and camera"
-        N = image.shape[0]
-
-        # encode image
-        image_feats = self.encoder(image)
-        assert image_feats.shape[-1] == self.encoder_feat_dim, \
-            f"Feature dimension mismatch: {image_feats.shape[-1]} vs {self.encoder_feat_dim}"
-
-        # embed camera
-        camera_embeddings = self.camera_embedder(camera)
-        assert camera_embeddings.shape[-1] == self.camera_embed_dim, \
-            f"Feature dimension mismatch: {camera_embeddings.shape[-1]} vs {self.camera_embed_dim}"
-
-        # transformer generating planes
-        planes = self.transformer(image_feats, camera_embeddings)
-        assert planes.shape[0] == N, "Batch size mismatch for planes"
-        assert planes.shape[1] == 3, "Planes should have 3 channels"
-        return planes
-

lrm/models/rendering/__init__.py

@@ -1,5 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.

lrm/models/rendering/synthesizer_part.py

@@ -1,194 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-
-import itertools
-import torch
-import torch.nn as nn
-
-from .utils.renderer import ImportanceRenderer
-from .utils.ray_sampler_part import RaySampler
-
-
-class OSGDecoder(nn.Module):
-    """
-    Triplane decoder that gives RGB and sigma values from sampled features.
-    Using ReLU here instead of Softplus in the original implementation.
-    
-    Reference:
-    EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L112
-    """
-    def __init__(self, n_features: int,
-                 hidden_dim: int = 64, num_layers: int = 4, activation: nn.Module = nn.ReLU):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(3 * n_features, hidden_dim),
-            activation(),
-            *itertools.chain(*[[
-                nn.Linear(hidden_dim, hidden_dim),
-                activation(),
-            ] for _ in range(num_layers - 2)]),
-            nn.Linear(hidden_dim, 1 + 3),
-        )
-        # init all bias to zero
-        for m in self.modules():
-            if isinstance(m, nn.Linear):
-                nn.init.zeros_(m.bias)
-
-    def forward(self, sampled_features, ray_directions):
-        # Aggregate features by mean
-        # sampled_features = sampled_features.mean(1)
-        # Aggregate features by concatenation
-        _N, n_planes, _M, _C = sampled_features.shape
-        sampled_features = sampled_features.permute(0, 2, 1, 3).reshape(_N, _M, n_planes*_C)
-        x = sampled_features
-
-        N, M, C = x.shape
-        x = x.contiguous().view(N*M, C)
-
-        x = self.net(x)
-        x = x.view(N, M, -1)
-        rgb = torch.sigmoid(x[..., 1:])*(1 + 2*0.001) - 0.001  # Uses sigmoid clamping from MipNeRF
-        sigma = x[..., 0:1]
-
-        return {'rgb': rgb, 'sigma': sigma}
-
-
-class TriplaneSynthesizer(nn.Module):
-    """
-    Synthesizer that renders a triplane volume with planes and a camera.
-    
-    Reference:
-    EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L19
-    """
-
-    DEFAULT_RENDERING_KWARGS = {
-        'ray_start': 'auto',
-        'ray_end': 'auto',
-        'box_warp': 2.,
-        'white_back': True,
-        'disparity_space_sampling': False,
-        'clamp_mode': 'softplus',
-        'sampler_bbox_min': -1.,
-        'sampler_bbox_max': 1.,
-    }
-
-    def __init__(self, triplane_dim: int, samples_per_ray: int):
-        super().__init__()
-
-        # attributes
-        self.triplane_dim = triplane_dim
-        self.rendering_kwargs = {
-            **self.DEFAULT_RENDERING_KWARGS,
-            'depth_resolution': samples_per_ray // 2,
-            'depth_resolution_importance': samples_per_ray // 2,
-        }
-
-        # renderings
-        self.renderer = ImportanceRenderer()
-        self.ray_sampler = RaySampler()
-
-        # modules
-        self.decoder = OSGDecoder(n_features=triplane_dim)
-
-    def forward(self, planes, cameras, render_size: int, crop_size: int, start_x: int, start_y:int):
-        # planes: (N, 3, D', H', W')
-        # cameras: (N, M, D_cam)
-        # render_size: int
-        assert planes.shape[0] == cameras.shape[0], "Batch size mismatch for planes and cameras"
-        N, M = cameras.shape[:2]
-        cam2world_matrix = cameras[..., :16].view(N, M, 4, 4)
-        intrinsics = cameras[..., 16:25].view(N, M, 3, 3)
-
-        # Create a batch of rays for volume rendering
-        ray_origins, ray_directions = self.ray_sampler(
-            cam2world_matrix=cam2world_matrix.reshape(-1, 4, 4),
-            intrinsics=intrinsics.reshape(-1, 3, 3),
-            render_size=render_size,
-            crop_size = crop_size,
-            start_x = start_x,
-            start_y = start_y
-        )
-        assert N*M == ray_origins.shape[0], "Batch size mismatch for ray_origins"
-        assert ray_origins.dim() == 3, "ray_origins should be 3-dimensional"
-        # Perform volume rendering
-        rgb_samples, depth_samples, weights_samples = self.renderer(
-            planes.repeat_interleave(M, dim=0), self.decoder, ray_origins, ray_directions, self.rendering_kwargs,
-        )
-
-        # Reshape into 'raw' neural-rendered image
-        Himg = Wimg = crop_size
-        rgb_images = rgb_samples.permute(0, 2, 1).reshape(N, M, rgb_samples.shape[-1], Himg, Wimg).contiguous()
-        depth_images = depth_samples.permute(0, 2, 1).reshape(N, M, 1, Himg, Wimg)
-        weight_images = weights_samples.permute(0, 2, 1).reshape(N, M, 1, Himg, Wimg)
-
-        return {
-            'images_rgb': rgb_images,
-            'images_depth': depth_images,
-            'images_weight': weight_images,
-        }
-
-    def forward_grid(self, planes, grid_size: int, aabb: torch.Tensor = None):
-        # planes: (N, 3, D', H', W')
-        # grid_size: int
-        # aabb: (N, 2, 3)
-        if aabb is None:
-            aabb = torch.tensor([
-                [self.rendering_kwargs['sampler_bbox_min']] * 3,
-                [self.rendering_kwargs['sampler_bbox_max']] * 3,
-            ], device=planes.device, dtype=planes.dtype).unsqueeze(0).repeat(planes.shape[0], 1, 1)
-        assert planes.shape[0] == aabb.shape[0], "Batch size mismatch for planes and aabb"
-        N = planes.shape[0]
-
-        # create grid points for triplane query
-        grid_points = []
-        for i in range(N):
-            grid_points.append(torch.stack(torch.meshgrid(
-                torch.linspace(aabb[i, 0, 0], aabb[i, 1, 0], grid_size, device=planes.device),
-                torch.linspace(aabb[i, 0, 1], aabb[i, 1, 1], grid_size, device=planes.device),
-                torch.linspace(aabb[i, 0, 2], aabb[i, 1, 2], grid_size, device=planes.device),
-                indexing='ij',
-            ), dim=-1).reshape(-1, 3))
-        cube_grid = torch.stack(grid_points, dim=0).to(planes.device)
-
-        features = self.forward_points(planes, cube_grid)
-
-        # reshape into grid
-        features = {
-            k: v.reshape(N, grid_size, grid_size, grid_size, -1)
-            for k, v in features.items()
-        }
-        return features
-
-    def forward_points(self, planes, points: torch.Tensor, chunk_size: int = 2**20):
-        # planes: (N, 3, D', H', W')
-        # points: (N, P, 3)
-        N, P = points.shape[:2]
-
-        # query triplane in chunks
-        outs = []
-        for i in range(0, points.shape[1], chunk_size):
-            chunk_points = points[:, i:i+chunk_size]
-
-            # query triplane
-            chunk_out = self.renderer.run_model_activated(
-                planes=planes,
-                decoder=self.decoder,
-                sample_coordinates=chunk_points,
-                sample_directions=torch.zeros_like(chunk_points),
-                options=self.rendering_kwargs,
-            )
-            outs.append(chunk_out)
-
-        # concatenate the outputs
-        point_features = {
-            k: torch.cat([out[k] for out in outs], dim=1)
-            for k in outs[0].keys()
-        }
-        
-        sig = point_features['sigma']
-        print(sig.mean(), sig.max(), sig.min())
-        return point_features

lrm/models/rendering/utils/__init__.py

@@ -1,14 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
-#
-# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
-# property and proprietary rights in and to this material, related
-# documentation and any modifications thereto. Any use, reproduction,
-# disclosure or distribution of this material and related documentation
-# without an express license agreement from NVIDIA CORPORATION or
-# its affiliates is strictly prohibited.

lrm/models/rendering/utils/math_utils.py

@@ -1,123 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-# MIT License
-
-# Copyright (c) 2022 Petr Kellnhofer
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-import torch
-
-def transform_vectors(matrix: torch.Tensor, vectors4: torch.Tensor) -> torch.Tensor:
-    """
-    Left-multiplies MxM @ NxM. Returns NxM.
-    """
-    res = torch.matmul(vectors4, matrix.T)
-    return res
-
-
-def normalize_vecs(vectors: torch.Tensor) -> torch.Tensor:
-    """
-    Normalize vector lengths.
-    """
-    return vectors / (torch.norm(vectors, dim=-1, keepdim=True))
-
-def torch_dot(x: torch.Tensor, y: torch.Tensor):
-    """
-    Dot product of two tensors.
-    """
-    return (x * y).sum(-1)
-
-
-def get_ray_limits_box(rays_o: torch.Tensor, rays_d: torch.Tensor, box_side_length):
-    """
-    Author: Petr Kellnhofer
-    Intersects rays with the [-1, 1] NDC volume.
-    Returns min and max distance of entry.
-    Returns -1 for no intersection.
-    https://www.scratchapixel.com/lessons/3d-basic-rendering/minimal-ray-tracer-rendering-simple-shapes/ray-box-intersection
-    """
-    o_shape = rays_o.shape
-    rays_o = rays_o.detach().reshape(-1, 3)
-    rays_d = rays_d.detach().reshape(-1, 3)
-
-
-    bb_min = [-1*(box_side_length/2), -1*(box_side_length/2), -1*(box_side_length/2)]
-    bb_max = [1*(box_side_length/2), 1*(box_side_length/2), 1*(box_side_length/2)]
-    bounds = torch.tensor([bb_min, bb_max], dtype=rays_o.dtype, device=rays_o.device)
-    is_valid = torch.ones(rays_o.shape[:-1], dtype=bool, device=rays_o.device)
-
-    # Precompute inverse for stability.
-    invdir = 1 / rays_d
-    sign = (invdir < 0).long()
-
-    # Intersect with YZ plane.
-    tmin = (bounds.index_select(0, sign[..., 0])[..., 0] - rays_o[..., 0]) * invdir[..., 0]
-    tmax = (bounds.index_select(0, 1 - sign[..., 0])[..., 0] - rays_o[..., 0]) * invdir[..., 0]
-
-    # Intersect with XZ plane.
-    tymin = (bounds.index_select(0, sign[..., 1])[..., 1] - rays_o[..., 1]) * invdir[..., 1]
-    tymax = (bounds.index_select(0, 1 - sign[..., 1])[..., 1] - rays_o[..., 1]) * invdir[..., 1]
-
-    # Resolve parallel rays.
-    is_valid[torch.logical_or(tmin > tymax, tymin > tmax)] = False
-
-    # Use the shortest intersection.
-    tmin = torch.max(tmin, tymin)
-    tmax = torch.min(tmax, tymax)
-
-    # Intersect with XY plane.
-    tzmin = (bounds.index_select(0, sign[..., 2])[..., 2] - rays_o[..., 2]) * invdir[..., 2]
-    tzmax = (bounds.index_select(0, 1 - sign[..., 2])[..., 2] - rays_o[..., 2]) * invdir[..., 2]
-
-    # Resolve parallel rays.
-    is_valid[torch.logical_or(tmin > tzmax, tzmin > tmax)] = False
-
-    # Use the shortest intersection.
-    tmin = torch.max(tmin, tzmin)
-    tmax = torch.min(tmax, tzmax)
-
-    # Mark invalid.
-    tmin[torch.logical_not(is_valid)] = -1
-    tmax[torch.logical_not(is_valid)] = -2
-
-    return tmin.reshape(*o_shape[:-1], 1), tmax.reshape(*o_shape[:-1], 1)
-
-
-def linspace(start: torch.Tensor, stop: torch.Tensor, num: int):
-    """
-    Creates a tensor of shape [num, *start.shape] whose values are evenly spaced from start to end, inclusive.
-    Replicates but the multi-dimensional bahaviour of numpy.linspace in PyTorch.
-    """
-    # create a tensor of 'num' steps from 0 to 1
-    steps = torch.arange(num, dtype=torch.float32, device=start.device) / (num - 1)
-
-    # reshape the 'steps' tensor to [-1, *([1]*start.ndim)] to allow for broadcastings
-    # - using 'steps.reshape([-1, *([1]*start.ndim)])' would be nice here but torchscript
-    #   "cannot statically infer the expected size of a list in this contex", hence the code below
-    for i in range(start.ndim):
-        steps = steps.unsqueeze(-1)
-
-    # the output starts at 'start' and increments until 'stop' in each dimension
-    out = start[None] + steps * (stop - start)[None]
-
-    return out

lrm/models/rendering/utils/ray_marcher.py

@@ -1,73 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
-#
-# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
-# property and proprietary rights in and to this material, related
-# documentation and any modifications thereto. Any use, reproduction,
-# disclosure or distribution of this material and related documentation
-# without an express license agreement from NVIDIA CORPORATION or
-# its affiliates is strictly prohibited.
-#
-# Modified by Zexin He
-# The modifications are subject to the same license as the original.
-
-
-"""
-The ray marcher takes the raw output of the implicit representation and uses the volume rendering equation to produce composited colors and depths.
-Based off of the implementation in MipNeRF (this one doesn't do any cone tracing though!)
-"""
-
-import torch
-import torch.nn as nn
-
-
-class MipRayMarcher2(nn.Module):
-    def __init__(self, activation_factory):
-        super().__init__()
-        self.activation_factory = activation_factory
-
-    def run_forward(self, colors, densities, depths, rendering_options):
-        
-        deltas = depths[:, :, 1:] - depths[:, :, :-1] 
-        colors_mid = (colors[:, :, :-1] + colors[:, :, 1:]) / 2
-        densities_mid = (densities[:, :, :-1] + densities[:, :, 1:]) / 2
-        depths_mid = (depths[:, :, :-1] + depths[:, :, 1:]) / 2
-
-    
-
-        # using factory mode for better usability
-        densities_mid = self.activation_factory(rendering_options)(densities_mid)
-
-        density_delta = densities_mid * deltas
-
-        alpha = 1 - torch.exp(-density_delta)
-
-        alpha_shifted = torch.cat([torch.ones_like(alpha[:, :, :1]), 1-alpha + 1e-10], -2)
-        weights = alpha * torch.cumprod(alpha_shifted, -2)[:, :, :-1]
-
-        composite_rgb = torch.sum(weights * colors_mid, -2)
-        weight_total = weights.sum(2)
-        composite_depth = torch.sum(weights * depths_mid, -2) / weight_total
-
-        # clip the composite to min/max range of depths
-        composite_depth = torch.nan_to_num(composite_depth, float('inf'))
-        composite_depth = torch.clamp(composite_depth, torch.min(depths), torch.max(depths))
-
-        if rendering_options.get('white_back', False):
-            composite_rgb = composite_rgb + 1 - weight_total
-
-        # rendered value scale is 0-1, comment out original mipnerf scaling
-        # composite_rgb = composite_rgb * 2 - 1 # Scale to (-1, 1)
-
-        return composite_rgb, composite_depth, weights
-
-
-    def forward(self, colors, densities, depths, rendering_options):
-        composite_rgb, composite_depth, weights = self.run_forward(colors, densities, depths, rendering_options)
-
-        return composite_rgb, composite_depth, weights

lrm/models/rendering/utils/ray_sampler_part.py

@@ -1,94 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
-#
-# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
-# property and proprietary rights in and to this material, related
-# documentation and any modifications thereto. Any use, reproduction,
-# disclosure or distribution of this material and related documentation
-# without an express license agreement from NVIDIA CORPORATION or
-# its affiliates is strictly prohibited.
-#
-# Modified by Zexin He
-# The modifications are subject to the same license as the original.
-
-
-"""
-The ray sampler is a module that takes in camera matrices and resolution and batches of rays.
-Expects cam2world matrices that use the OpenCV camera coordinate system conventions.
-"""
-
-import torch
-
-class RaySampler(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.ray_origins_h, self.ray_directions, self.depths, self.image_coords, self.rendering_options = None, None, None, None, None
-
-
-    def forward(self, cam2world_matrix, intrinsics, render_size, crop_size, start_x, start_y):
-        """
-        Create batches of rays and return origins and directions.
-
-        cam2world_matrix: (N, 4, 4)
-        intrinsics: (N, 3, 3)
-        render_size: int
-
-        ray_origins: (N, M, 3)
-        ray_dirs: (N, M, 2)
-        """
-
-        N, M = cam2world_matrix.shape[0], crop_size**2
-        cam_locs_world = cam2world_matrix[:, :3, 3]
-        fx = intrinsics[:, 0, 0]
-        fy = intrinsics[:, 1, 1]
-        cx = intrinsics[:, 0, 2]
-        cy = intrinsics[:, 1, 2]
-        sk = intrinsics[:, 0, 1]
-
-        uv = torch.stack(torch.meshgrid(
-            torch.arange(render_size, dtype=torch.float32, device=cam2world_matrix.device),
-            torch.arange(render_size, dtype=torch.float32, device=cam2world_matrix.device),
-            indexing='ij',
-        )) 
-        if crop_size < render_size:
-            patch_uv = []
-            for i in range(cam2world_matrix.shape[0]):
-                patch_uv.append(uv.clone()[None, :, start_y:start_y+crop_size, start_x:start_x+crop_size])
-            uv = torch.cat(patch_uv, 0)
-            uv = uv.flip(1).reshape(cam2world_matrix.shape[0], 2, -1).transpose(2, 1)
-        else:
-            uv = uv.flip(0).reshape(2, -1).transpose(1, 0)
-            uv = uv.unsqueeze(0).repeat(cam2world_matrix.shape[0], 1, 1)
-        # uv = uv.unsqueeze(0).repeat(cam2world_matrix.shape[0], 1, 1)
-        # uv = uv.flip(1).reshape(cam2world_matrix.shape[0], 2, -1).transpose(2, 1)
-        x_cam = uv[:, :, 0].view(N, -1) * (1./render_size) + (0.5/render_size)
-        y_cam = uv[:, :, 1].view(N, -1) * (1./render_size) + (0.5/render_size)
-        z_cam = torch.ones((N, M), device=cam2world_matrix.device)
-
-        x_lift = (x_cam - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y_cam/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z_cam
-        y_lift = (y_cam - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z_cam
-
-        cam_rel_points = torch.stack((x_lift, y_lift, z_cam, torch.ones_like(z_cam)), dim=-1).float()
-
-        _opencv2blender = torch.tensor([
-            [1, 0, 0, 0],
-            [0, -1, 0, 0],
-            [0, 0, -1, 0],
-            [0, 0, 0, 1],
-        ], dtype=torch.float32, device=cam2world_matrix.device).unsqueeze(0).repeat(N, 1, 1)
-
-        # added float here
-        cam2world_matrix = torch.bmm(cam2world_matrix.float(), _opencv2blender.float())
-
-        world_rel_points = torch.bmm(cam2world_matrix.float(), cam_rel_points.permute(0, 2, 1)).permute(0, 2, 1)[:, :, :3]
-
-        ray_dirs = world_rel_points - cam_locs_world[:, None, :]
-        ray_dirs = torch.nn.functional.normalize(ray_dirs, dim=2)
-
-        ray_origins = cam_locs_world.unsqueeze(1).repeat(1, ray_dirs.shape[1], 1)
-        return ray_origins, ray_dirs

lrm/models/rendering/utils/renderer.py

@@ -1,314 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
-#
-# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
-# property and proprietary rights in and to this material, related
-# documentation and any modifications thereto. Any use, reproduction,
-# disclosure or distribution of this material and related documentation
-# without an express license agreement from NVIDIA CORPORATION or
-# its affiliates is strictly prohibited.
-#
-# Modified by Zexin He
-# The modifications are subject to the same license as the original.
-
-
-"""
-The renderer is a module that takes in rays, decides where to sample along each
-ray, and computes pixel colors using the volume rendering equation.
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .ray_marcher import MipRayMarcher2
-from . import math_utils
-
-def generate_planes():
-    """
-    Defines planes by the three vectors that form the "axes" of the
-    plane. Should work with arbitrary number of planes and planes of
-    arbitrary orientation.
-
-    Bugfix reference: https://github.com/NVlabs/eg3d/issues/67
-    """
-    return torch.tensor([[[1, 0, 0],
-                            [0, 1, 0],
-                            [0, 0, 1]],
-                            [[1, 0, 0],
-                            [0, 0, 1],
-                            [0, 1, 0]],
-                            [[0, 0, 1],
-                            [0, 1, 0],
-                            [1, 0, 0]]], dtype=torch.float32)
-
-def project_onto_planes(planes, coordinates):
-    """
-    Does a projection of a 3D point onto a batch of 2D planes,
-    returning 2D plane coordinates.
-
-    Takes plane axes of shape n_planes, 3, 3
-    # Takes coordinates of shape N, M, 3
-    # returns projections of shape N*n_planes, M, 2
-    """
-    N, M, C = coordinates.shape
-    n_planes, _, _ = planes.shape
-    coordinates = coordinates.unsqueeze(1).expand(-1, n_planes, -1, -1).reshape(N*n_planes, M, 3)
-    inv_planes = torch.linalg.inv(planes).unsqueeze(0).expand(N, -1, -1, -1).reshape(N*n_planes, 3, 3)
-    coordinates = coordinates.to(inv_planes.device)
-    projections = torch.bmm(coordinates, inv_planes)
-    return projections[..., :2]
-
-def sample_from_planes(plane_axes, plane_features, coordinates, mode='bilinear', padding_mode='zeros', box_warp=None):
-    assert padding_mode == 'zeros'
-    N, n_planes, C, H, W = plane_features.shape
-    _, M, _ = coordinates.shape
-    plane_features = plane_features.view(N*n_planes, C, H, W)
-
-    coordinates = (2/box_warp) * coordinates # add specific box bounds
-    # half added here
-    projected_coordinates = project_onto_planes(plane_axes, coordinates).unsqueeze(1)
-    # removed float from projected_coordinates
-    output_features = torch.nn.functional.grid_sample(plane_features.float(), projected_coordinates.float(), mode=mode, padding_mode=padding_mode, align_corners=False).permute(0, 3, 2, 1).reshape(N, n_planes, M, C)
-    return output_features
-
-def sample_from_3dgrid(grid, coordinates):
-    """
-    Expects coordinates in shape (batch_size, num_points_per_batch, 3)
-    Expects grid in shape (1, channels, H, W, D)
-    (Also works if grid has batch size)
-    Returns sampled features of shape (batch_size, num_points_per_batch, feature_channels)
-    """
-    batch_size, n_coords, n_dims = coordinates.shape
-    sampled_features = torch.nn.functional.grid_sample(grid.expand(batch_size, -1, -1, -1, -1),
-                                                       coordinates.reshape(batch_size, 1, 1, -1, n_dims),
-                                                       mode='bilinear', padding_mode='zeros', align_corners=False)
-    N, C, H, W, D = sampled_features.shape
-    sampled_features = sampled_features.permute(0, 4, 3, 2, 1).reshape(N, H*W*D, C)
-    return sampled_features
-
-class ImportanceRenderer(torch.nn.Module):
-    """
-    Modified original version to filter out-of-box samples as TensoRF does.
-    
-    Reference:
-    TensoRF: https://github.com/apchenstu/TensoRF/blob/main/models/tensorBase.py#L277
-    """
-    def __init__(self):
-        super().__init__()
-        self.activation_factory = self._build_activation_factory()
-        self.ray_marcher = MipRayMarcher2(self.activation_factory)
-        self.plane_axes = generate_planes()
-
-    def _build_activation_factory(self):
-        def activation_factory(options: dict):
-            if options['clamp_mode'] == 'softplus':
-                return lambda x: F.softplus(x - 1)  # activation bias of -1 makes things initialize better
-            else:
-                assert False, "Renderer only supports `clamp_mode`=`softplus`!"
-        return activation_factory
-
-    def _forward_pass(self, depths: torch.Tensor, ray_directions: torch.Tensor, ray_origins: torch.Tensor,
-                        planes: torch.Tensor, decoder: nn.Module, rendering_options: dict):
-        """
-        Additional filtering is applied to filter out-of-box samples.
-        Modifications made by Zexin He.
-        """
-
-        # context related variables
-        batch_size, num_rays, samples_per_ray, _ = depths.shape
-        device = planes.device
-        depths = depths.to(device)
-        ray_directions = ray_directions.to(device)
-        ray_origins = ray_origins.to(device)
-        # define sample points with depths
-        sample_directions = ray_directions.unsqueeze(-2).expand(-1, -1, samples_per_ray, -1).reshape(batch_size, -1, 3)
-        sample_coordinates = (ray_origins.unsqueeze(-2) + depths * ray_directions.unsqueeze(-2)).reshape(batch_size, -1, 3)
-
-        # filter out-of-box samples
-        mask_inbox = \
-            (rendering_options['sampler_bbox_min'] <= sample_coordinates) & \
-                (sample_coordinates <= rendering_options['sampler_bbox_max'])
-        mask_inbox = mask_inbox.all(-1)
-
-        # forward model according to all samples
-        _out = self.run_model(planes, decoder, sample_coordinates, sample_directions, rendering_options)
-
-        # set out-of-box samples to zeros(rgb) & -inf(sigma)
-        SAFE_GUARD = 3
-        DATA_TYPE = _out['sigma'].dtype
-        colors_pass = torch.zeros(batch_size, num_rays * samples_per_ray, 3, device=device, dtype=DATA_TYPE)
-        densities_pass = torch.nan_to_num(torch.full((batch_size, num_rays * samples_per_ray, 1), -float('inf'), device=device, dtype=DATA_TYPE)) / SAFE_GUARD
-        colors_pass[mask_inbox], densities_pass[mask_inbox] = _out['rgb'][mask_inbox], _out['sigma'][mask_inbox]
-
-        # reshape back
-        colors_pass = colors_pass.reshape(batch_size, num_rays, samples_per_ray, colors_pass.shape[-1])
-        densities_pass = densities_pass.reshape(batch_size, num_rays, samples_per_ray, densities_pass.shape[-1])
-
-        return colors_pass, densities_pass
-
-    def forward(self, planes, decoder, ray_origins, ray_directions, rendering_options):
-        # self.plane_axes = self.plane_axes.to(ray_origins.device)
-
-        if rendering_options['ray_start'] == rendering_options['ray_end'] == 'auto':
-            ray_start, ray_end = math_utils.get_ray_limits_box(ray_origins, ray_directions, box_side_length=rendering_options['box_warp'])
-            is_ray_valid = ray_end > ray_start
-            if torch.any(is_ray_valid).item():
-                ray_start[~is_ray_valid] = ray_start[is_ray_valid].min()
-                ray_end[~is_ray_valid] = ray_start[is_ray_valid].max()
-            depths_coarse = self.sample_stratified(ray_origins, ray_start, ray_end, rendering_options['depth_resolution'], rendering_options['disparity_space_sampling'])
-        else:
-            # Create stratified depth samples
-            depths_coarse = self.sample_stratified(ray_origins, rendering_options['ray_start'], rendering_options['ray_end'], rendering_options['depth_resolution'], rendering_options['disparity_space_sampling'])
-        
-        depths_coarse =  depths_coarse.to(planes.device)
-
-        # Coarse Pass
-        colors_coarse, densities_coarse = self._forward_pass(
-            depths=depths_coarse, ray_directions=ray_directions, ray_origins=ray_origins,
-            planes=planes, decoder=decoder, rendering_options=rendering_options)
-
-        # Fine Pass
-        N_importance = rendering_options['depth_resolution_importance']
-        if N_importance > 0:
-            _, _, weights = self.ray_marcher(colors_coarse, densities_coarse, depths_coarse, rendering_options)
-
-            depths_fine = self.sample_importance(depths_coarse, weights, N_importance)
-
-            colors_fine, densities_fine = self._forward_pass(
-                depths=depths_fine, ray_directions=ray_directions, ray_origins=ray_origins,
-                planes=planes, decoder=decoder, rendering_options=rendering_options)
-
-            all_depths, all_colors, all_densities = self.unify_samples(depths_coarse, colors_coarse, densities_coarse,
-                                                                  depths_fine, colors_fine, densities_fine)
-
-            # Aggregate
-            rgb_final, depth_final, weights = self.ray_marcher(all_colors, all_densities, all_depths, rendering_options)
-        else:
-            rgb_final, depth_final, weights = self.ray_marcher(colors_coarse, densities_coarse, depths_coarse, rendering_options)
-
-        return rgb_final, depth_final, weights.sum(2)
-
-    def run_model(self, planes, decoder, sample_coordinates, sample_directions, options):
-        plane_axes = self.plane_axes.to(planes.device)
-        sampled_features = sample_from_planes(plane_axes, planes, sample_coordinates, padding_mode='zeros', box_warp=options['box_warp'])
-
-        out = decoder(sampled_features, sample_directions)
-        if options.get('density_noise', 0) > 0:
-            out['sigma'] += torch.randn_like(out['sigma']) * options['density_noise']
-        return out
-
-    def run_model_activated(self, planes, decoder, sample_coordinates, sample_directions, options):
-        out = self.run_model(planes, decoder, sample_coordinates, sample_directions, options)
-        out['sigma'] = self.activation_factory(options)(out['sigma'])
-        return out
-
-    def sort_samples(self, all_depths, all_colors, all_densities):
-        _, indices = torch.sort(all_depths, dim=-2)
-        all_depths = torch.gather(all_depths, -2, indices)
-        all_colors = torch.gather(all_colors, -2, indices.expand(-1, -1, -1, all_colors.shape[-1]))
-        all_densities = torch.gather(all_densities, -2, indices.expand(-1, -1, -1, 1))
-        return all_depths, all_colors, all_densities
-
-    def unify_samples(self, depths1, colors1, densities1, depths2, colors2, densities2):
-        all_depths = torch.cat([depths1, depths2], dim = -2)
-        all_colors = torch.cat([colors1, colors2], dim = -2)
-        all_densities = torch.cat([densities1, densities2], dim = -2)
-
-        _, indices = torch.sort(all_depths, dim=-2)
-        all_depths = torch.gather(all_depths, -2, indices)
-        all_colors = torch.gather(all_colors, -2, indices.expand(-1, -1, -1, all_colors.shape[-1]))
-        all_densities = torch.gather(all_densities, -2, indices.expand(-1, -1, -1, 1))
-
-        return all_depths, all_colors, all_densities
-
-    def sample_stratified(self, ray_origins, ray_start, ray_end, depth_resolution, disparity_space_sampling=False):
-        """
-        Return depths of approximately uniformly spaced samples along rays.
-        """
-        N, M, _ = ray_origins.shape
-        if disparity_space_sampling:
-            depths_coarse = torch.linspace(0,
-                                    1,
-                                    depth_resolution,
-                                    device=ray_origins.device).reshape(1, 1, depth_resolution, 1).repeat(N, M, 1, 1)
-            depth_delta = 1/(depth_resolution - 1)
-            depths_coarse += torch.rand_like(depths_coarse) * depth_delta
-            depths_coarse = 1./(1./ray_start * (1. - depths_coarse) + 1./ray_end * depths_coarse)
-        else:
-            if type(ray_start) == torch.Tensor:
-                depths_coarse = math_utils.linspace(ray_start, ray_end, depth_resolution).permute(1,2,0,3)
-                depth_delta = (ray_end - ray_start) / (depth_resolution - 1)
-                depths_coarse += torch.rand_like(depths_coarse) * depth_delta[..., None]
-            else:
-                depths_coarse = torch.linspace(ray_start, ray_end, depth_resolution, device=ray_origins.device).reshape(1, 1, depth_resolution, 1).repeat(N, M, 1, 1)
-                depth_delta = (ray_end - ray_start)/(depth_resolution - 1)
-                depths_coarse += torch.rand_like(depths_coarse) * depth_delta
-
-        return depths_coarse
-
-    def sample_importance(self, z_vals, weights, N_importance):
-        """
-        Return depths of importance sampled points along rays. See NeRF importance sampling for more.
-        """
-        with torch.no_grad():
-            batch_size, num_rays, samples_per_ray, _ = z_vals.shape
-
-            z_vals = z_vals.reshape(batch_size * num_rays, samples_per_ray)
-            weights = weights.reshape(batch_size * num_rays, -1) # -1 to account for loss of 1 sample in MipRayMarcher
-
-            # smooth weights
-            weights = torch.nn.functional.max_pool1d(weights.unsqueeze(1).float(), 2, 1, padding=1)
-            weights = torch.nn.functional.avg_pool1d(weights, 2, 1).squeeze()
-            weights = weights + 0.01
-
-            z_vals_mid = 0.5 * (z_vals[: ,:-1] + z_vals[: ,1:])
-            importance_z_vals = self.sample_pdf(z_vals_mid, weights[:, 1:-1],
-                                             N_importance).detach().reshape(batch_size, num_rays, N_importance, 1)
-        return importance_z_vals
-
-    def sample_pdf(self, bins, weights, N_importance, det=False, eps=1e-5):
-        """
-        Sample @N_importance samples from @bins with distribution defined by @weights.
-        Inputs:
-            bins: (N_rays, N_samples_+1) where N_samples_ is "the number of coarse samples per ray - 2"
-            weights: (N_rays, N_samples_)
-            N_importance: the number of samples to draw from the distribution
-            det: deterministic or not
-            eps: a small number to prevent division by zero
-        Outputs:
-            samples: the sampled samples
-        """
-        N_rays, N_samples_ = weights.shape
-        weights = weights + eps # prevent division by zero (don't do inplace op!)
-        pdf = weights / torch.sum(weights, -1, keepdim=True) # (N_rays, N_samples_)
-        cdf = torch.cumsum(pdf, -1) # (N_rays, N_samples), cumulative distribution function
-        cdf = torch.cat([torch.zeros_like(cdf[: ,:1]), cdf], -1)  # (N_rays, N_samples_+1)
-                                                                   # padded to 0~1 inclusive
-
-        if det:
-            u = torch.linspace(0, 1, N_importance, device=bins.device)
-            u = u.expand(N_rays, N_importance)
-        else:
-            u = torch.rand(N_rays, N_importance, device=bins.device)
-        u = u.contiguous()
-
-        inds = torch.searchsorted(cdf, u, right=True)
-        below = torch.clamp_min(inds-1, 0)
-        above = torch.clamp_max(inds, N_samples_)
-
-        inds_sampled = torch.stack([below, above], -1).view(N_rays, 2*N_importance)
-        cdf_g = torch.gather(cdf, 1, inds_sampled).view(N_rays, N_importance, 2)
-        bins_g = torch.gather(bins, 1, inds_sampled).view(N_rays, N_importance, 2)
-
-        denom = cdf_g[...,1]-cdf_g[...,0]
-        denom[denom<eps] = 1 # denom equals 0 means a bin has weight 0, in which case it will not be sampled
-                             # anyway, therefore any value for it is fine (set to 1 here)
-
-        samples = bins_g[...,0] + (u-cdf_g[...,0])/denom * (bins_g[...,1]-bins_g[...,0])
-        return samples

lrm/models/transformer.py

@@ -1,135 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-
-import torch
-import torch.nn as nn
-
-
-class ModLN(nn.Module):
-    """
-    Modulation with adaLN.
-    
-    References:
-    DiT: https://github.com/facebookresearch/DiT/blob/main/models.py#L101
-    """
-    def __init__(self, inner_dim: int, mod_dim: int, eps: float):
-        super().__init__()
-        self.norm = nn.LayerNorm(inner_dim, eps=eps)
-        self.mlp = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(mod_dim, inner_dim * 2),
-        )
-
-    @staticmethod
-    def modulate(x, shift, scale):
-        # x: [N, L, D]
-        # shift, scale: [N, D]
-        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-    def forward(self, x, cond):
-        shift, scale = self.mlp(cond).chunk(2, dim=-1)  # [N, D]
-        return self.modulate(self.norm(x), shift, scale)  # [N, L, D]
-
-
-class ConditionModulationBlock(nn.Module):
-    """
-    Transformer block that takes in a cross-attention condition and another modulation vector applied to sub-blocks.
-    """
-    # use attention from torch.nn.MultiHeadAttention
-    # Block contains a cross-attention layer, a self-attention layer, and a MLP
-    def __init__(self, inner_dim: int, cond_dim: int, mod_dim: int, num_heads: int, eps: float,
-                 attn_drop: float = 0., attn_bias: bool = False,
-                 mlp_ratio: float = 4., mlp_drop: float = 0.):
-        super().__init__()
-        self.norm1 = ModLN(inner_dim, mod_dim, eps)
-        self.cross_attn = nn.MultiheadAttention(
-            embed_dim=inner_dim, num_heads=num_heads, kdim=cond_dim, vdim=cond_dim,
-            dropout=attn_drop, bias=attn_bias, batch_first=True)
-        self.norm2 = ModLN(inner_dim, mod_dim, eps)
-        self.self_attn = nn.MultiheadAttention(
-            embed_dim=inner_dim, num_heads=num_heads,
-            dropout=attn_drop, bias=attn_bias, batch_first=True)
-        self.norm3 = ModLN(inner_dim, mod_dim, eps)
-        self.mlp = nn.Sequential(
-            nn.Linear(inner_dim, int(inner_dim * mlp_ratio)),
-            nn.GELU(),
-            nn.Dropout(mlp_drop),
-            nn.Linear(int(inner_dim * mlp_ratio), inner_dim),
-            nn.Dropout(mlp_drop),
-        )
-
-    def forward(self, x, cond, mod):
-        # x: [N, L, D]
-        # cond: [N, L_cond, D_cond]
-        # mod: [N, D_mod]
-        x = x + self.cross_attn(self.norm1(x, mod), cond, cond, need_weights=False)[0]
-        before_sa = self.norm2(x, mod)
-        x = x + self.self_attn(before_sa, before_sa, before_sa, need_weights=False)[0]
-        x = x + self.mlp(self.norm3(x, mod))
-        return x
-
-
-class TriplaneTransformer(nn.Module):
-    """
-    Transformer with condition and modulation that generates a triplane representation.
-    
-    Reference:
-    Timm: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L486
-    """
-    def __init__(self, inner_dim: int, image_feat_dim: int, camera_embed_dim: int,
-                 triplane_low_res: int, triplane_high_res: int, triplane_dim: int,
-                 num_layers: int, num_heads: int,
-                 eps: float = 1e-6):
-        super().__init__()
-
-        # attributes
-        self.triplane_low_res = triplane_low_res
-        self.triplane_high_res = triplane_high_res
-        self.triplane_dim = triplane_dim
-
-        # modules
-        # initialize pos_embed with 1/sqrt(dim) * N(0, 1)
-        self.pos_embed = nn.Parameter(torch.randn(1, 3*triplane_low_res**2, inner_dim) * (1. / inner_dim) ** 0.5)
-        self.layers = nn.ModuleList([
-            ConditionModulationBlock(
-                inner_dim=inner_dim, cond_dim=image_feat_dim, mod_dim=camera_embed_dim, num_heads=num_heads, eps=eps)
-            for _ in range(num_layers)
-        ])
-        self.norm = nn.LayerNorm(inner_dim, eps=eps)
-        self.deconv = nn.ConvTranspose2d(inner_dim, triplane_dim, kernel_size=2, stride=2, padding=0)
-
-    def forward(self, image_feats, camera_embeddings):
-        # image_feats: [N, L_cond, D_cond]
-        # camera_embeddings: [N, D_mod]
-
-        assert image_feats.shape[0] == camera_embeddings.shape[0], \
-            f"Mismatched batch size: {image_feats.shape[0]} vs {camera_embeddings.shape[0]}"
-
-        N = image_feats.shape[0]
-        H = W = self.triplane_low_res
-        L = 3 * H * W
-
-        x = self.pos_embed.repeat(N, 1, 1)  # [N, L, D]
-        for layer in self.layers:
-            x = layer(x, image_feats, camera_embeddings)
-        x = self.norm(x)
-
-        # separate each plane and apply deconv
-        x = x.view(N, 3, H, W, -1)
-        x = torch.einsum('nihwd->indhw', x)  # [3, N, D, H, W]
-        x = x.contiguous().view(3*N, -1, H, W)  # [3*N, D, H, W]
-        x = self.deconv(x)  # [3*N, D', H', W']
-        x = x.view(3, N, *x.shape[-3:])  # [3, N, D', H', W']
-        x = torch.einsum('indhw->nidhw', x)  # [N, 3, D', H', W']
-        x = x.contiguous()
-
-        assert self.triplane_high_res == x.shape[-2], \
-            f"Output triplane resolution does not match with expected: {x.shape[-2]} vs {self.triplane_high_res}"
-        assert self.triplane_dim == x.shape[-3], \
-            f"Output triplane dimension does not match with expected: {x.shape[-3]} vs {self.triplane_dim}"
-
-        return x