openlrm/runners/infer/lrm.py

# Copyright (c) 2023-2024, Zexin He
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import torch
import os
import argparse
import mcubes
import trimesh
import safetensors
import numpy as np
from PIL import Image
from omegaconf import OmegaConf
from tqdm.auto import tqdm
from accelerate.logging import get_logger
from huggingface_hub import hf_hub_download

from .base_inferrer import Inferrer
from openlrm.datasets.cam_utils import build_camera_principle, build_camera_standard, surrounding_views_linspace, create_intrinsics
from openlrm.utils.logging import configure_logger
from openlrm.runners import REGISTRY_RUNNERS
from openlrm.utils.video import images_to_video
from openlrm.utils.hf_hub import wrap_model_hub


logger = get_logger(__name__)


def parse_configs():

    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str)
    parser.add_argument('--infer', type=str)
    args, unknown = parser.parse_known_args()

    cfg = OmegaConf.create()
    cli_cfg = OmegaConf.from_cli(unknown)

    # parse from ENV
    if os.environ.get('APP_INFER') is not None:
        args.infer = os.environ.get('APP_INFER')
    if os.environ.get('APP_MODEL_NAME') is not None:
        cli_cfg.model_name = os.environ.get('APP_MODEL_NAME')
    if os.environ.get('APP_PRETRAIN_MODEL_NAME') is not None:
        cli_cfg.pretrain_model_hf = os.environ.get('APP_PRETRAIN_MODEL_NAME')

    if args.config is not None:
        cfg_train = OmegaConf.load(args.config)
        cfg.source_size = cfg_train.dataset.source_image_res
        cfg.render_size = cfg_train.dataset.render_image.high
        _relative_path = os.path.join(cfg_train.experiment.parent, cfg_train.experiment.child, os.path.basename(cli_cfg.model_name).split('_')[-1])
        cfg.video_dump = os.path.join("exps", 'videos', _relative_path)
        cfg.mesh_dump = os.path.join("exps", 'meshes', _relative_path)

    if args.infer is not None:
        cfg_infer = OmegaConf.load(args.infer)
        cfg.merge_with(cfg_infer)
        if hasattr(cfg, 'experiment') and hasattr(cfg.experiment, 'parent'):
            cfg.setdefault('video_dump', os.path.join("dumps", cli_cfg.model_name, cfg.experiment.parent, cfg.experiment.child, 'videos'))
            cfg.setdefault('mesh_dump', os.path.join("dumps", cli_cfg.model_name, cfg.experiment.parent, cfg.experiment.child, 'meshes'))
        else:
            cfg.setdefault('video_dump', os.path.join("dumps", cli_cfg.model_name, 'videos'))
            cfg.setdefault('mesh_dump', os.path.join("dumps", cli_cfg.model_name, 'meshes'))
            
    cfg.setdefault('double_sided', False)
    cfg.setdefault('pretrain_model_hf', None)
    cfg.merge_with(cli_cfg)

    """
    [required]
    model_name: str
    image_input: str
    export_video: bool
    export_mesh: bool

    [special]
    source_size: int
    render_size: int
    video_dump: str
    mesh_dump: str

    [default]
    render_views: int
    render_fps: int
    mesh_size: int
    mesh_thres: float
    frame_size: int
    logger: str
    """

    cfg.setdefault('inferrer', {})
    cfg['inferrer'].setdefault('logger', 'INFO')

    # assert not (args.config is not None and args.infer is not None), "Only one of config and infer should be provided"
    assert cfg.model_name is not None, "model_name is required"
    if not os.environ.get('APP_ENABLED', None):
        assert cfg.image_input is not None, "image_input is required"
        assert cfg.export_video or cfg.export_mesh, \
            "At least one of export_video or export_mesh should be True"
        cfg.app_enabled = False
    else:
        cfg.app_enabled = True

    return cfg


@REGISTRY_RUNNERS.register('infer.lrm')
class LRMInferrer(Inferrer):

    EXP_TYPE: str = 'lrm'

    def __init__(self):
        super().__init__()

        self.cfg = parse_configs()
        configure_logger(
            stream_level=self.cfg.inferrer.logger,
            log_level=self.cfg.inferrer.logger,
        )

        self.model = self._build_model(self.cfg).to(self.device)

    def _load_checkpoint(self, cfg):
        ckpt_root = os.path.join(
            cfg.saver.checkpoint_root,
            cfg.experiment.parent, cfg.experiment.child,
        )
        if not os.path.exists(ckpt_root):
            raise FileNotFoundError(f"The checkpoint directory '{ckpt_root}' does not exist.")
        ckpt_dirs = os.listdir(ckpt_root)
        iter_number = "{:06}".format(cfg.inferrer.iteration)
        if iter_number not in ckpt_dirs:
            raise FileNotFoundError(f"Checkpoint for iteration '{iter_number}' not found in '{ckpt_root}'.")
        inferrer_ckpt_path = os.path.join(ckpt_root, iter_number, 'model.safetensors')
        logger.info(f"======== Auto-resume from {inferrer_ckpt_path} ========")
        return inferrer_ckpt_path

    def _build_model(self, cfg):
        from openlrm.models import model_dict
        if cfg.inferrer.hugging_face is True:    # for huggingface infer
            hf_model_cls = wrap_model_hub(model_dict[self.EXP_TYPE])
            model = hf_model_cls.from_pretrained(cfg.model_name)
            if cfg.double_sided:
                pretrain_model_path = hf_hub_download(repo_id=cfg.pretrain_model_hf, filename='model.safetensors')
                safetensors.torch.load_model(       # load the pretrain model after load the Tailor3D finetune part.
                    model,
                    pretrain_model_path,
                    strict=False
                )
        else: # for common infer
            model = model_dict[self.EXP_TYPE](**cfg['model'])
            inferrer_ckpt_path = self._load_checkpoint(cfg)
            if cfg.double_sided:
                pretrain_model_path = hf_hub_download(repo_id=cfg.pretrain_model_hf, filename='model.safetensors')
                safetensors.torch.load_model(       # load the pretrain model.
                    model,
                    pretrain_model_path,
                    strict=False
                )
                safetensors.torch.load_model(       # load the finetune model.
                    model,
                    inferrer_ckpt_path,
                    strict=False
                )
            else:
                safetensors.torch.load_model(
                    model,
                    inferrer_ckpt_path,
                )                
        return model
    
    @staticmethod
    def save_images(images, output_path):
        os.makedirs((output_path), exist_ok=True)
        for i in range(images.shape[0]):
            frame = (images[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
            Image.fromarray(frame).save(os.path.join(output_path, f"{str(i)}.png"))

    def _default_source_camera(self, dist_to_center: float = 2.0, batch_size: int = 1, device: torch.device = torch.device('cpu')):
        # return: (N, D_cam_raw)
        canonical_camera_extrinsics = torch.tensor([[
            [1, 0, 0, 0],
            [0, 0, -1, -dist_to_center],
            [0, 1, 0, 0],
        ]], dtype=torch.float32, device=device)
        canonical_camera_intrinsics = create_intrinsics(
            f=0.75,
            c=0.5,
            device=device,
        ).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, n_views: int, batch_size: int = 1, device: torch.device = torch.device('cpu')):
        # return: (N, M, D_cam_render)
        render_camera_extrinsics = surrounding_views_linspace(n_views=n_views, device=device)
        render_camera_intrinsics = create_intrinsics(
            f=0.75,
            c=0.5,
            device=device,
        ).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)

    def infer_planes(self, image: torch.Tensor, source_cam_dist: float, back_image=None):
        N = image.shape[0]
        source_camera = self._default_source_camera(dist_to_center=source_cam_dist, batch_size=N, device=self.device)
        front_planes = self.model.forward_planes(image, source_camera)
        if back_image is not None:
            back_planes = self.model.forward_planes(back_image, source_camera)
            # XY Plane
            back_planes[:, 0, :, :, :] = torch.flip(back_planes[:, 0, :, :, :], dims=[-2, -1])
            # XZ Plane
            back_planes[:, 1, :, :, :] = torch.flip(back_planes[:, 1, :, :, :], dims=[-1])
            # YZ Plane
            back_planes[:, 2, :, :, :] = torch.flip(back_planes[:, 2, :, :, :], dims=[-2])

            # To fuse the front planes and the back planes
            bs, num_planes, channels, height, width = front_planes.shape
            if 'conv_fuse' in self.cfg['model']:
                planes = torch.cat((front_planes, back_planes), dim=2)
                planes = planes.reshape(-1, channels*2, height, width) 
                # planes = self.model.front_back_conv(planes).view(bs, num_planes, -1, height, width)  # only one layer.
                # Apply multiple convolutional layers
                for layer in self.model.front_back_conv:
                    planes = layer(planes)
                
                planes = planes.view(bs, num_planes, -1, height, width)
            elif 'swin_ca_fuse' in self.cfg['model']:
                front_planes = front_planes.reshape(bs*num_planes, channels, height*width).permute(0, 2, 1).contiguous()    # [8, 3, 32, 64, 64] -> [24, 32, 4096] -> [24, 4096, 32]
                back_planes = back_planes.reshape(bs*num_planes, channels, height*width).permute(0, 2, 1).contiguous()
                planes = self.model.swin_cross_attention(front_planes, back_planes, height, width)[0].permute(0, 2, 1).reshape(bs, num_planes, channels, height, width)
        else:
            planes = front_planes

        assert N == planes.shape[0]
        return planes

    def infer_video(self, planes: torch.Tensor, frame_size: int, render_size: int, render_views: int, render_fps: int, dump_video_path: str, image_format=False):
        N = planes.shape[0]
        render_cameras = self._default_render_cameras(n_views=render_views, batch_size=N, device=self.device)
        render_anchors = torch.zeros(N, render_cameras.shape[1], 2, device=self.device)
        render_resolutions = torch.ones(N, render_cameras.shape[1], 1, device=self.device) * render_size
        render_bg_colors = torch.ones(N, render_cameras.shape[1], 1, device=self.device, dtype=torch.float32) * 1.

        frames = []
        for i in range(0, render_cameras.shape[1], frame_size):
            frames.append(
                self.model.synthesizer(
                    planes=planes,
                    cameras=render_cameras[:, i:i+frame_size],
                    anchors=render_anchors[:, i:i+frame_size],
                    resolutions=render_resolutions[:, i:i+frame_size],
                    bg_colors=render_bg_colors[:, i:i+frame_size],
                    region_size=render_size,
                )
            )
        # merge frames
        frames = {
            k: torch.cat([r[k] for r in frames], dim=1)
            for k in frames[0].keys()
        }
        # dump
        os.makedirs(os.path.dirname(dump_video_path), exist_ok=True)
        for k, v in frames.items():
            if k == 'images_rgb':
                if image_format:
                    self.save_images(                                       # save the rendering images directly.
                        v[0],
                        os.path.join(dump_video_path.replace('.mov', ''), 'nvs'),
                    )
                else:
                    images_to_video(
                        images=v[0],
                        output_path=dump_video_path,
                        fps=render_fps,
                        gradio_codec=self.cfg.app_enabled,
                    )

    def infer_mesh(self, planes: torch.Tensor, mesh_size: int, mesh_thres: float, dump_mesh_path: str):
        grid_out = self.model.synthesizer.forward_grid(
            planes=planes,
            grid_size=mesh_size,
        )
        
        vtx, faces = mcubes.marching_cubes(grid_out['sigma'].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'].squeeze(0).cpu().numpy()  # (0, 1)
        vtx_colors = (vtx_colors * 255).astype(np.uint8)
        
        mesh = trimesh.Trimesh(vertices=vtx, faces=faces, vertex_colors=vtx_colors)

        # dump
        os.makedirs(os.path.dirname(dump_mesh_path), exist_ok=True)
        mesh.export(dump_mesh_path)

    def infer_single(self, image_path: str, source_cam_dist: float, export_video: bool, export_mesh: bool, dump_video_path: str, dump_mesh_path: str, image_path_back=None):
        source_size = self.cfg.inferrer.source_size
        render_size = self.cfg.inferrer.render_size
        render_views = self.cfg.inferrer.render_views
        render_fps = self.cfg.inferrer.render_fps
        mesh_size = self.cfg.inferrer.mesh_size
        mesh_thres = self.cfg.inferrer.mesh_thres
        frame_size = self.cfg.inferrer.frame_size
        source_cam_dist = self.cfg.inferrer.source_cam_dist if source_cam_dist is None else source_cam_dist

        image_format = self.cfg.inferrer.image_format

        image = self.open_image(image_path, source_size)
        if image_path_back is None:
            back_image = self.open_image(image_path.replace('front', 'back'), source_size) if self.cfg.double_sided else None
        else:
            back_image = self.open_image(image_path_back, source_size) if self.cfg.double_sided else None

        with torch.no_grad():
            planes = self.infer_planes(image, source_cam_dist=source_cam_dist, back_image=back_image)

            results = {}
            if export_video:
                frames = self.infer_video(planes, frame_size=frame_size, render_size=render_size, render_views=render_views, render_fps=render_fps, dump_video_path=dump_video_path, 
                                          image_format=image_format)
                results.update({
                    'frames': frames,
                })
            if export_mesh:
                mesh = self.infer_mesh(planes, mesh_size=mesh_size, mesh_thres=mesh_thres, dump_mesh_path=dump_mesh_path)
                results.update({
                    'mesh': mesh,
                })

    def data_init(self):
        image_paths = []
        if os.path.isfile(self.cfg.image_input):
            omit_prefix = os.path.dirname(self.cfg.image_input)
            image_paths.append(self.cfg.image_input)
        else:
            omit_prefix = self.cfg.image_input
            if self.cfg.double_sided:                   # double sided
                walk_path = os.path.join(self.cfg.image_input, 'front')
            else:
                walk_path = self.cfg.image_input
            for root, dirs, files in os.walk(walk_path):
                for file in files:
                    if file.endswith('.png'):
                        image_paths.append(os.path.join(root, file))
            image_paths.sort()
        # alloc to each DDP worker
        image_paths = image_paths[self.accelerator.process_index::self.accelerator.num_processes]

        return image_paths, omit_prefix

    def open_image(self, image_path, source_size):
        # prepare image: [1, C_img, H_img, W_img], 0-1 scale
        image = torch.from_numpy(np.array(Image.open(image_path))).to(self.device)
        image = image.permute(2, 0, 1).unsqueeze(0) / 255.0
        if image.shape[1] == 4:  # RGBA
            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)

        return image

    def infer(self):
        image_paths, omit_prefix = self.data_init()
        for image_path in tqdm(image_paths, disable=not self.accelerator.is_local_main_process):

            # prepare dump paths
            image_name = os.path.basename(image_path)
            uid = image_name.split('.')[0]
            subdir_path = os.path.dirname(image_path).replace(omit_prefix, '')
            subdir_path = subdir_path[1:] if subdir_path.startswith('/') else subdir_path
            dump_video_path = os.path.join(
                self.cfg.video_dump,
                subdir_path,
                f'{uid}.mov',
            )
            dump_mesh_path = os.path.join(
                self.cfg.mesh_dump,
                subdir_path,
                f'{uid}.ply',
            )

            self.infer_single(
                image_path,
                source_cam_dist=None,
                export_video=self.cfg.export_video,
                export_mesh=self.cfg.export_mesh,
                dump_video_path=dump_video_path,
                dump_mesh_path=dump_mesh_path,
            )