modeling.py

#### modeling.py
import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig
import torch
import numpy as np
import math

from .dino_wrapper2 import DinoWrapper
from .transformer import TriplaneTransformer
from .synthesizer_part import TriplaneSynthesizer

class CameraEmbedder(nn.Module):
    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 LRMGeneratorConfig(PretrainedConfig):
    model_type = "lrm_generator"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.camera_embed_dim = kwargs.get("camera_embed_dim", 1024)
        self.rendering_samples_per_ray = kwargs.get("rendering_samples_per_ray", 128)
        self.transformer_dim = kwargs.get("transformer_dim", 1024)
        self.transformer_layers = kwargs.get("transformer_layers", 16)
        self.transformer_heads = kwargs.get("transformer_heads", 16)
        self.triplane_low_res = kwargs.get("triplane_low_res", 32)
        self.triplane_high_res = kwargs.get("triplane_high_res", 64)
        self.triplane_dim = kwargs.get("triplane_dim", 80)
        self.encoder_freeze = kwargs.get("encoder_freeze", False)
        self.encoder_model_name = kwargs.get("encoder_model_name", 'facebook/dinov2-base')
        self.encoder_feat_dim = kwargs.get("encoder_feat_dim", 768)

    

class LRMGenerator(PreTrainedModel):
    config_class = LRMGeneratorConfig

    def __init__(self, config: LRMGeneratorConfig):
        super().__init__(config)

        self.encoder_feat_dim = config.encoder_feat_dim
        self.camera_embed_dim = config.camera_embed_dim

        self.encoder = DinoWrapper(
            model_name=config.encoder_model_name,
            freeze=config.encoder_freeze,
        )
        self.camera_embedder = CameraEmbedder(
            raw_dim=12 + 4, embed_dim=config.camera_embed_dim,
        )
        self.transformer = TriplaneTransformer(
            inner_dim=config.transformer_dim, num_layers=config.transformer_layers, num_heads=config.transformer_heads,
            image_feat_dim=config.encoder_feat_dim,
            camera_embed_dim=config.camera_embed_dim,
            triplane_low_res=config.triplane_low_res, triplane_high_res=config.triplane_high_res, triplane_dim=config.triplane_dim,
        )
        self.synthesizer = TriplaneSynthesizer(
            triplane_dim=config.triplane_dim, samples_per_ray=config.rendering_samples_per_ray,
        )

    def forward(self, image, camera, export_mesh=False, mesh_size=512, render_size=384, export_video=False, fps=30):

        assert image.shape[0] == camera.shape[0], "Batch size mismatch"
        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}"

        with torch.no_grad():

            # 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"
    
            # Generate the mesh
            if export_mesh:
              import mcubes
              import trimesh
              grid_out = self.synthesizer.forward_grid(planes=planes, grid_size=mesh_size)
              vtx, faces = mcubes.marching_cubes(grid_out['sigma'].float().squeeze(0).squeeze(-1).cpu().numpy(), 1.0)
              vtx = vtx / (mesh_size - 1) * 2 - 1
              vtx_tensor = torch.tensor(vtx, dtype=torch.float32, device=image.device).unsqueeze(0)
              vtx_colors = self.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)
            
              mesh_path = "awesome_mesh.obj"
              mesh.export(mesh_path, 'obj')
    
              return planes, mesh_path
            
            # Generate video
            if export_video:
                render_cameras = self._default_render_cameras(batch_size=N).to(image.device)
    
                frames = []
                chunk_size = 1  # Adjust chunk size as needed
                for i in range(0, render_cameras.shape[1], chunk_size):
                    frame_chunk = self.synthesizer(
                        planes,
                        render_cameras[:, i:i + chunk_size],
                        render_size,
                        render_size,
                        0,
                        0
                    )
                    frames.append(frame_chunk['images_rgb'])
    
                frames = torch.cat(frames, dim=1)
                frames = (frames.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8)
    
                # Save video
                video_path = "awesome_video.mp4"
                imageio.mimwrite(video_path, frames, fps=fps)
    
                return planes, video_path
    
            return planes
        
    # Copied from https://github.com/facebookresearch/vfusion3d/blob/main/lrm/cam_utils.py
    # and https://github.com/facebookresearch/vfusion3d/blob/main/lrm/inferrer.py
    def _default_intrinsics(self):
        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_render_cameras(self, batch_size=1):
        M = 160  # Number of views
        radius = 1.5
        elevation = 0

        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 = self.center_looking_at_camera_pose(camera_positions)

        intrinsics = self._default_intrinsics().unsqueeze(0).repeat(extrinsics.shape[0], 1, 1)
        render_cameras = self.build_camera_standard(extrinsics, intrinsics)

        return render_cameras.unsqueeze(0).repeat(batch_size, 1, 1)

    def center_looking_at_camera_pose(self, camera_position: torch.Tensor, look_at: torch.Tensor = None, up_world: torch.Tensor = None):
        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_normalized_camera_intrinsics(self, intrinsics: torch.Tensor):
        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_standard(self, RT: torch.Tensor, intrinsics: torch.Tensor):
        E = self.compose_extrinsic_RT(RT)
        fx, fy, cx, cy = self.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 compose_extrinsic_RT(self, RT: torch.Tensor):
        return torch.cat([
            RT,
            torch.tensor([[[0, 0, 0, 1]]], dtype=torch.float32).repeat(RT.shape[0], 1, 1).to(RT.device)
        ], dim=1)