benchmarks.wiki / Public workspace
LongBench v2 / 66ebd3ba5a08c7b9b35e0446 / The Instant3D paper introduced significant innovations in accelerating 3D object…
Problem
Answer published by the source. Consult the official source to check your work against its answer.
question
The Instant3D paper introduced significant innovations in accelerating 3D object generation by reducing the computational complexity typically seen in traditional 3D modeling methods. In contrast, OpenLRM adopts a large-scale reconstruction approach that leverages hybrid datasets like Objaverse and MVImgNet. Considering these two frameworks, how might OpenLRM’s use of hybrid datasets and a large reconstruction model introduce new challenges that were less emphasized in the Instant3D approach, especially in terms of real-time object generation and scaling?
context · full text (453,905 characters)
# OpenLRM: Open-Source Large Reconstruction Models
[](LICENSE)
[](LICENSE_WEIGHT)
[](https://arxiv.org/abs/2311.04400)
[](https://huggingface.co/zxhezexin)
[](https://huggingface.co/spaces/zxhezexin/OpenLRM)
<img src="assets/rendered_video/teaser.gif" width="75%" height="auto"/>
<div style="text-align: left">
<img src="assets/mesh_snapshot/crop.owl.ply00.png" width="12%" height="auto"/>
<img src="assets/mesh_snapshot/crop.owl.ply01.png" width="12%" height="auto"/>
<img src="assets/mesh_snapshot/crop.building.ply00.png" width="12%" height="auto"/>
<img src="assets/mesh_snapshot/crop.building.ply01.png" width="12%" height="auto"/>
<img src="assets/mesh_snapshot/crop.rose.ply00.png" width="12%" height="auto"/>
<img src="assets/mesh_snapshot/crop.rose.ply01.png" width="12%" height="auto"/>
</div>
## News
- [2024.03.13] Update [training code](openlrm/runners/train) and release [OpenLRM v1.1.1](https://github.com/3DTopia/OpenLRM/releases/tag/v1.1.1).
- [2024.03.08] We have released the core [blender script](scripts/data/objaverse/blender_script.py) used to render Objaverse images.
- [2024.03.05] The [Huggingface demo](https://huggingface.co/spaces/zxhezexin/OpenLRM) now uses `openlrm-mix-base-1.1` model by default. Please refer to the [model card](model_card.md) for details on the updated model architecture and training settings.
- [2024.03.04] Version update v1.1. Release model weights trained on both Objaverse and MVImgNet. Codebase is majorly refactored for better usability and extensibility. Please refer to [v1.1.0](https://github.com/3DTopia/OpenLRM/releases/tag/v1.1.0) for details.
- [2024.01.09] Updated all v1.0 models trained on Objaverse. Please refer to [HF Models](https://huggingface.co/zxhezexin) and overwrite previous model weights.
- [2023.12.21] [Hugging Face Demo](https://huggingface.co/spaces/zxhezexin/OpenLRM) is online. Have a try!
- [2023.12.20] Release weights of the base and large models trained on Objaverse.
- [2023.12.20] We release this project OpenLRM, which is an open-source implementation of the paper [LRM](https://arxiv.org/abs/2311.04400).
## Setup
### Installation
Code
git clone https://github.com/3DTopia/OpenLRM.git
cd OpenLRM
### Environment - Install requirements for OpenLRM first.
Code
pip install -r requirements.txt
- Please then follow the [xFormers installation guide](https://github.com/facebookresearch/xformers?tab=readme-ov-file#installing-xformers) to enable memory efficient attention inside [DINOv2 encoder](openlrm/models/encoders/dinov2/layers/attention.py). ## Quick Start ### Pretrained Models - Model weights are released on [Hugging Face](https://huggingface.co/zxhezexin). - Weights will be downloaded automatically when you run the inference script for the first time. - Please be aware of the [license](LICENSE_WEIGHT) before using the weights. | Model | Training Data | Layers | Feat. Dim | Trip. Dim. | In. Res. | Link | | :--- | :--- | :--- | :--- | :--- | :--- | :--- | | openlrm-obj-small-1.1 | Objaverse | 12 | 512 | 32 | 224 | [HF](https://huggingface.co/zxhezexin/openlrm-obj-small-1.1) | | openlrm-obj-base-1.1 | Objaverse | 12 | 768 | 48 | 336 | [HF](https://huggingface.co/zxhezexin/openlrm-obj-base-1.1) | | openlrm-obj-large-1.1 | Objaverse | 16 | 1024 | 80 | 448 | [HF](https://huggingface.co/zxhezexin/openlrm-obj-large-1.1) | | openlrm-mix-small-1.1 | Objaverse + MVImgNet | 12 | 512 | 32 | 224 | [HF](https://huggingface.co/zxhezexin/openlrm-mix-small-1.1) | | openlrm-mix-base-1.1 | Objaverse + MVImgNet | 12 | 768 | 48 | 336 | [HF](https://huggingface.co/zxhezexin/openlrm-mix-base-1.1) | | openlrm-mix-large-1.1 | Objaverse + MVImgNet | 16 | 1024 | 80 | 448 | [HF](https://huggingface.co/zxhezexin/openlrm-mix-large-1.1) | Model cards with additional details can be found in [model_card.md](model_card.md). ### Prepare Images - We put some sample inputs under `assets/sample_input`, and you can quickly try them. - Prepare RGBA images or RGB images with white background (with some background removal tools, e.g., [Rembg](https://github.com/danielgatis/rembg), [Clipdrop](https://clipdrop.co)). ### Inference - Run the inference script to get 3D assets. - You may specify which form of output to generate by setting the flags `EXPORT_VIDEO=true` and `EXPORT_MESH=true`. - Please set default `INFER_CONFIG` according to the model you want to use. E.g., `infer-b.yaml` for base models and `infer-s.yaml` for small models. - An example usage is as follows:
Code
# Example usage
EXPORT_VIDEO=true
EXPORT_MESH=true
INFER_CONFIG="./configs/infer-b.yaml"
MODEL_NAME="zxhezexin/openlrm-mix-base-1.1"
IMAGE_INPUT="./assets/sample_input/owl.png"
python -m openlrm.launch infer.lrm --infer $INFER_CONFIG model_name=$MODEL_NAME image_input=$IMAGE_INPUT export_video=$EXPORT_VIDEO export_mesh=$EXPORT_MESH
### Tips - The recommended PyTorch version is `>=2.1`. Code is developed and tested under PyTorch `2.1.2`. - If you encounter CUDA OOM issues, please try to reduce the `frame_size` in the inference configs. - You should be able to see `UserWarning: xFormers is available` if `xFormers` is actually working. ## Training ### Configuration - We provide a sample accelerate config file under `configs/accelerate-train.yaml`, which defaults to use 8 GPUs with `bf16` mixed precision. - You may modify the configuration file to fit your own environment. ### Data Preparation - We provide the core [Blender script](scripts/data/objaverse/blender_script.py) used to render Objaverse images. - Please refer to [Objaverse Rendering](https://github.com/allenai/objaverse-rendering) for other scripts including distributed rendering. ### Run Training - A sample training config file is provided under `configs/train-sample.yaml`. - Please replace data related paths in the config file with your own paths and customize the training settings. - An example training usage is as follows:
Code
# Example usage
ACC_CONFIG="./configs/accelerate-train.yaml"
TRAIN_CONFIG="./configs/train-sample.yaml"
accelerate launch --config_file $ACC_CONFIG -m openlrm.launch train.lrm --config $TRAIN_CONFIG
### Inference on Trained Models - The inference pipeline is compatible with huggingface utilities for better convenience. - You need to convert the training checkpoint to inference models by running the following script.
Code
python scripts/convert_hf.py --config <YOUR_EXACT_TRAINING_CONFIG> convert.global_step=null
- The converted model will be saved under `exps/releases` by default and can be used for inference following the [inference guide](https://github.com/3DTopia/OpenLRM?tab=readme-ov-file#inference). ## Acknowledgement - We thank the authors of the [original paper](https://arxiv.org/abs/2311.04400) for their great work! Special thanks to Kai Zhang and Yicong Hong for assistance during the reproduction. - This project is supported by Shanghai AI Lab by providing the computing resources. - This project is advised by Ziwei Liu and Jiaya Jia. ## Citation If you find this work useful for your research, please consider citing:
Code
@article{hong2023lrm,
title={Lrm: Large reconstruction model for single image to 3d},
author={Hong, Yicong and Zhang, Kai and Gu, Jiuxiang and Bi, Sai and Zhou, Yang and Liu, Difan and Liu, Feng and Sunkavalli, Kalyan and Bui, Trung and Tan, Hao},
journal={arXiv preprint arXiv:2311.04400},
year={2023}
}
Code
@misc{openlrm,
title = {OpenLRM: Open-Source Large Reconstruction Models},
author = {Zexin He and Tengfei Wang},
year = {2023},
howpublished = {\url{https://github.com/3DTopia/OpenLRM}},
}
## License
- OpenLRM as a whole is licensed under the [Apache License, Version 2.0](LICENSE), while certain components are covered by [NVIDIA's proprietary license](LICENSE_NVIDIA). Users are responsible for complying with the respective licensing terms of each component.
- Model weights are licensed under the [Creative Commons Attribution-NonCommercial 4.0 International License](LICENSE_WEIGHT). They are provided for research purposes only, and CANNOT be used commercially.
# Model Card for OpenLRM V1.1
## Overview
- This model card is for the [OpenLRM](https://github.com/3DTopia/OpenLRM) project, which is an open-source implementation of the paper [LRM](https://arxiv.org/abs/2311.04400).
- Information contained in this model card corresponds to [Version 1.1](https://github.com/3DTopia/OpenLRM/releases).
## Model Details
- Training data
| Model | Training Data |
| :---: | :---: |
| [openlrm-obj-small-1.1](https://huggingface.co/zxhezexin/openlrm-obj-small-1.1) | Objaverse |
| [openlrm-obj-base-1.1](https://huggingface.co/zxhezexin/openlrm-obj-base-1.1) | Objaverse |
| [openlrm-obj-large-1.1](https://huggingface.co/zxhezexin/openlrm-obj-large-1.1) | Objaverse |
| [openlrm-mix-small-1.1](https://huggingface.co/zxhezexin/openlrm-mix-small-1.1) | Objaverse + MVImgNet |
| [openlrm-mix-base-1.1](https://huggingface.co/zxhezexin/openlrm-mix-base-1.1) | Objaverse + MVImgNet |
| [openlrm-mix-large-1.1](https://huggingface.co/zxhezexin/openlrm-mix-large-1.1) | Objaverse + MVImgNet |
- Model architecture (version==1.1)
| Type | Layers | Feat. Dim | Attn. Heads | Triplane Dim. | Input Res. | Image Encoder | Encoder Dim. | Size |
| :---: | :----: | :-------: | :---------: | :-----------: | :--------: | :---------------: | :----------: | :---: |
| small | 12 | 512 | 8 | 32 | 224 | dinov2_vits14_reg | 384 | 446M |
| base | 12 | 768 | 12 | 48 | 336 | dinov2_vitb14_reg | 768 | 1.04G |
| large | 16 | 1024 | 16 | 80 | 448 | dinov2_vitb14_reg | 768 | 1.81G |
- Training settings
| Type | Rend. Res. | Rend. Patch | Ray Samples |
| :---: | :--------: | :---------: | :---------: |
| small | 192 | 64 | 96 |
| base | 288 | 96 | 96 |
| large | 384 | 128 | 128 |
## Notable Differences from the Original Paper
- We do not use the deferred back-propagation technique in the original paper.
- We used random background colors during training.
- The image encoder is based on the [DINOv2](https://github.com/facebookresearch/dinov2) model with register tokens.
- The triplane decoder contains 4 layers in our implementation.
## License
- The model weights are released under the [Creative Commons Attribution-NonCommercial 4.0 International License](LICENSE_WEIGHT).
- They are provided for research purposes only, and CANNOT be used commercially.
## Disclaimer
This model is an open-source implementation and is NOT the official release of the original research paper. While it aims to reproduce the original results as faithfully as possible, there may be variations due to model implementation, training data, and other factors.
### Ethical Considerations
- This model should be used responsibly and ethically, and should not be used for malicious purposes.
- Users should be aware of potential biases in the training data.
- The model should not be used under the circumstances that could lead to harm or unfair treatment of individuals or groups.
### Usage Considerations
- The model is provided "as is" without warranty of any kind.
- Users are responsible for ensuring that their use complies with all relevant laws and regulations.
- The developers and contributors of this model are not liable for any damages or losses arising from the use of this model.
---
*This model card is subject to updates and modifications. Users are advised to check for the latest version regularly.*
torch>=2.1.2
lpips
omegaconf
transformers
safetensors
accelerate
imageio[ffmpeg]
PyMCubes
trimesh
megfile
opencv-python
optimum[onnxruntime-gpu]
rembg[gpu,cli]
httpx[socks]
tensorboard
# 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 os
from PIL import Image
import numpy as np
import gradio as gr
def assert_input_image(input_image):
if input_image is None:
raise gr.Error("No image selected or uploaded!")
def prepare_working_dir():
import tempfile
working_dir = tempfile.TemporaryDirectory()
return working_dir
def init_preprocessor():
from openlrm.utils.preprocess import Preprocessor
global preprocessor
preprocessor = Preprocessor()
def preprocess_fn(image_in: np.ndarray, remove_bg: bool, recenter: bool, working_dir):
image_raw = os.path.join(working_dir.name, "raw.png")
with Image.fromarray(image_in) as img:
img.save(image_raw)
image_out = os.path.join(working_dir.name, "rembg.png")
success = preprocessor.preprocess(image_path=image_raw, save_path=image_out, rmbg=remove_bg, recenter=recenter)
assert success, f"Failed under preprocess_fn!"
return image_out
def demo_openlrm(infer_impl):
def core_fn(image: str, source_cam_dist: float, working_dir):
dump_video_path = os.path.join(working_dir.name, "output.mp4")
dump_mesh_path = os.path.join(working_dir.name, "output.ply")
infer_impl(
image_path=image,
source_cam_dist=source_cam_dist,
export_video=True,
export_mesh=False,
dump_video_path=dump_video_path,
dump_mesh_path=dump_mesh_path,
)
return dump_video_path
def example_fn(image: np.ndarray):
from gradio.utils import get_cache_folder
working_dir = get_cache_folder()
image = preprocess_fn(
image_in=image,
remove_bg=True,
recenter=True,
working_dir=working_dir,
)
video = core_fn(
image=image,
source_cam_dist=2.0,
working_dir=working_dir,
)
return image, video
_TITLE = '''OpenLRM: Open-Source Large Reconstruction Models'''
_DESCRIPTION = '''
<div>
<a style="display:inline-block" href='https://github.com/3DTopia/OpenLRM'><img src='https://img.shields.io/github/stars/3DTopia/OpenLRM?style=social'/></a>
<a style="display:inline-block; margin-left: .5em" href="https://huggingface.co/zxhezexin"><img src='https://img.shields.io/badge/Model-Weights-blue'/></a>
</div>
OpenLRM is an open-source implementation of Large Reconstruction Models.
<strong>Image-to-3D in 10 seconds with A100!</strong>
<strong>Disclaimer:</strong> This demo uses `openlrm-mix-base-1.1` model with 288x288 rendering resolution here for a quick demonstration.
'''
with gr.Blocks(analytics_enabled=False) as demo:
# HEADERS
with gr.Row():
with gr.Column(scale=1):
gr.Markdown('# ' + _TITLE)
with gr.Row():
gr.Markdown(_DESCRIPTION)
# DISPLAY
with gr.Row():
with gr.Column(variant='panel', scale=1):
with gr.Tabs(elem_id="openlrm_input_image"):
with gr.TabItem('Input Image'):
with gr.Row():
input_image = gr.Image(label="Input Image", image_mode="RGBA", width="auto", sources="upload", type="numpy", elem_id="content_image")
with gr.Column(variant='panel', scale=1):
with gr.Tabs(elem_id="openlrm_processed_image"):
with gr.TabItem('Processed Image'):
with gr.Row():
processed_image = gr.Image(label="Processed Image", image_mode="RGBA", type="filepath", elem_id="processed_image", width="auto", interactive=False)
with gr.Column(variant='panel', scale=1):
with gr.Tabs(elem_id="openlrm_render_video"):
with gr.TabItem('Rendered Video'):
with gr.Row():
output_video = gr.Video(label="Rendered Video", format="mp4", width="auto", autoplay=True)
# SETTING
with gr.Row():
with gr.Column(variant='panel', scale=1):
with gr.Tabs(elem_id="openlrm_attrs"):
with gr.TabItem('Settings'):
with gr.Column(variant='panel'):
gr.Markdown(
"""
<strong>Best Practice</strong>:
Centered objects in reasonable sizes. Try adjusting source camera distances.
"""
)
checkbox_rembg = gr.Checkbox(True, label='Remove background')
checkbox_recenter = gr.Checkbox(True, label='Recenter the object')
slider_cam_dist = gr.Slider(1.0, 3.5, value=2.0, step=0.1, label="Source Camera Distance")
submit = gr.Button('Generate', elem_id="openlrm_generate", variant='primary')
# EXAMPLES
with gr.Row():
examples = [
['assets/sample_input/owl.png'],
['assets/sample_input/building.png'],
['assets/sample_input/mailbox.png'],
['assets/sample_input/fire.png'],
['assets/sample_input/girl.png'],
['assets/sample_input/lamp.png'],
['assets/sample_input/hydrant.png'],
['assets/sample_input/hotdogs.png'],
['assets/sample_input/traffic.png'],
['assets/sample_input/ceramic.png'],
]
gr.Examples(
examples=examples,
inputs=[input_image],
outputs=[processed_image, output_video],
fn=example_fn,
cache_examples=bool(os.getenv('SPACE_ID')),
examples_per_page=20,
)
working_dir = gr.State()
submit.click(
fn=assert_input_image,
inputs=[input_image],
queue=False,
).success(
fn=prepare_working_dir,
outputs=[working_dir],
queue=False,
).success(
fn=preprocess_fn,
inputs=[input_image, checkbox_rembg, checkbox_recenter, working_dir],
outputs=[processed_image],
).success(
fn=core_fn,
inputs=[processed_image, slider_cam_dist, working_dir],
outputs=[output_video],
)
demo.queue()
demo.launch()
def launch_gradio_app():
os.environ.update({
"APP_ENABLED": "1",
"APP_MODEL_NAME": "zxhezexin/openlrm-mix-base-1.1",
"APP_INFER": "./configs/infer-gradio.yaml",
"APP_TYPE": "infer.lrm",
"NUMBA_THREADING_LAYER": 'omp',
})
from openlrm.runners import REGISTRY_RUNNERS
from openlrm.runners.infer.base_inferrer import Inferrer
InferrerClass : Inferrer = REGISTRY_RUNNERS[os.getenv("APP_TYPE")]
with InferrerClass() as inferrer:
init_preprocessor()
if not bool(os.getenv('SPACE_ID')):
from openlrm.utils.proxy import no_proxy
demo = no_proxy(demo_openlrm)
else:
demo = demo_openlrm
demo(infer_impl=inferrer.infer_single)
if __name__ == '__main__':
launch_gradio_app()
# 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 argparse
from openlrm.runners import REGISTRY_RUNNERS
def main():
parser = argparse.ArgumentParser(description='OpenLRM launcher')
parser.add_argument('runner', type=str, help='Runner to launch')
args, unknown = parser.parse_known_args()
if args.runner not in REGISTRY_RUNNERS:
raise ValueError('Runner {} not found'.format(args.runner))
RunnerClass = REGISTRY_RUNNERS[args.runner]
with RunnerClass() as runner:
runner.run()
if __name__ == '__main__':
main()
# 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.
#
# Empty
# 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.
from abc import ABC, abstractmethod
class Runner(ABC):
"""Abstract runner class"""
def __init__(self):
pass
@abstractmethod
def run(self):
pass
# 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.
from openlrm.utils.registry import Registry
REGISTRY_RUNNERS = Registry()
from .train import *
from .infer import *
# 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 os
import math
from tqdm.auto import tqdm
import torch
import torch.nn as nn
from torchvision.utils import make_grid
from accelerate.logging import get_logger
from .base_trainer import Trainer
from openlrm.utils.profiler import DummyProfiler
from openlrm.runners import REGISTRY_RUNNERS
logger = get_logger(__name__)
@REGISTRY_RUNNERS.register('train.lrm')
class LRMTrainer(Trainer):
def __init__(self):
super().__init__()
self.model = self._build_model(self.cfg)
self.optimizer = self._build_optimizer(self.model, self.cfg)
self.train_loader, self.val_loader = self._build_dataloader(self.cfg)
self.scheduler = self._build_scheduler(self.optimizer, self.cfg)
self.pixel_loss_fn, self.perceptual_loss_fn, self.tv_loss_fn = self._build_loss_fn(self.cfg)
def _build_model(self, cfg):
assert cfg.experiment.type == 'lrm', \
f"Config type {cfg.experiment.type} does not match with runner {self.__class__.__name__}"
from openlrm.models import ModelLRM
model = ModelLRM(**cfg.model)
return model
def _build_optimizer(self, model: nn.Module, cfg):
decay_params, no_decay_params = [], []
# add all bias and LayerNorm params to no_decay_params
for name, module in model.named_modules():
if isinstance(module, nn.LayerNorm):
no_decay_params.extend([p for p in module.parameters()])
elif hasattr(module, 'bias') and module.bias is not None:
no_decay_params.append(module.bias)
# add remaining parameters to decay_params
_no_decay_ids = set(map(id, no_decay_params))
decay_params = [p for p in model.parameters() if id(p) not in _no_decay_ids]
# filter out parameters with no grad
decay_params = list(filter(lambda p: p.requires_grad, decay_params))
no_decay_params = list(filter(lambda p: p.requires_grad, no_decay_params))
# monitor this to make sure we don't miss any parameters
logger.info("======== Weight Decay Parameters ========")
logger.info(f"Total: {len(decay_params)}")
logger.info("======== No Weight Decay Parameters ========")
logger.info(f"Total: {len(no_decay_params)}")
# Optimizer
opt_groups = [
{'params': decay_params, 'weight_decay': cfg.train.optim.weight_decay},
{'params': no_decay_params, 'weight_decay': 0.0},
]
optimizer = torch.optim.AdamW(
opt_groups,
lr=cfg.train.optim.lr,
betas=(cfg.train.optim.beta1, cfg.train.optim.beta2),
)
return optimizer
def _build_scheduler(self, optimizer, cfg):
local_batches_per_epoch = math.floor(len(self.train_loader) / self.accelerator.num_processes)
total_global_batches = cfg.train.epochs * math.ceil(local_batches_per_epoch / self.cfg.train.accum_steps)
effective_warmup_iters = cfg.train.scheduler.warmup_real_iters
logger.debug(f"======== Scheduler effective max iters: {total_global_batches} ========")
logger.debug(f"======== Scheduler effective warmup iters: {effective_warmup_iters} ========")
if cfg.train.scheduler.type == 'cosine':
from openlrm.utils.scheduler import CosineWarmupScheduler
scheduler = CosineWarmupScheduler(
optimizer=optimizer,
warmup_iters=effective_warmup_iters,
max_iters=total_global_batches,
)
else:
raise NotImplementedError(f"Scheduler type {cfg.train.scheduler.type} not implemented")
return scheduler
def _build_dataloader(self, cfg):
# dataset class
from openlrm.datasets import MixerDataset
# build dataset
train_dataset = MixerDataset(
split="train",
subsets=cfg.dataset.subsets,
sample_side_views=cfg.dataset.sample_side_views,
render_image_res_low=cfg.dataset.render_image.low,
render_image_res_high=cfg.dataset.render_image.high,
render_region_size=cfg.dataset.render_image.region,
source_image_res=cfg.dataset.source_image_res,
normalize_camera=cfg.dataset.normalize_camera,
normed_dist_to_center=cfg.dataset.normed_dist_to_center,
)
val_dataset = MixerDataset(
split="val",
subsets=cfg.dataset.subsets,
sample_side_views=cfg.dataset.sample_side_views,
render_image_res_low=cfg.dataset.render_image.low,
render_image_res_high=cfg.dataset.render_image.high,
render_region_size=cfg.dataset.render_image.region,
source_image_res=cfg.dataset.source_image_res,
normalize_camera=cfg.dataset.normalize_camera,
normed_dist_to_center=cfg.dataset.normed_dist_to_center,
)
# build data loader
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.train.batch_size,
shuffle=True,
drop_last=True,
num_workers=cfg.dataset.num_train_workers,
pin_memory=cfg.dataset.pin_mem,
persistent_workers=True,
)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=cfg.val.batch_size,
shuffle=False,
drop_last=False,
num_workers=cfg.dataset.num_val_workers,
pin_memory=cfg.dataset.pin_mem,
persistent_workers=False,
)
return train_loader, val_loader
def _build_loss_fn(self, cfg):
from openlrm.losses import PixelLoss, LPIPSLoss, TVLoss
pixel_loss_fn = PixelLoss()
with self.accelerator.main_process_first():
perceptual_loss_fn = LPIPSLoss(device=self.device, prefech=True)
tv_loss_fn = TVLoss()
return pixel_loss_fn, perceptual_loss_fn, tv_loss_fn
def register_hooks(self):
pass
def forward_loss_local_step(self, data):
source_camera = data['source_camera']
render_camera = data['render_camera']
source_image = data['source_image']
render_image = data['render_image']
render_anchors = data['render_anchors']
render_full_resolutions = data['render_full_resolutions']
render_bg_colors = data['render_bg_colors']
N, M, C, H, W = render_image.shape
# forward
outputs = self.model(
image=source_image,
source_camera=source_camera,
render_cameras=render_camera,
render_anchors=render_anchors,
render_resolutions=render_full_resolutions,
render_bg_colors=render_bg_colors,
render_region_size=self.cfg.dataset.render_image.region,
)
# loss calculation
loss = 0.
loss_pixel = None
loss_perceptual = None
loss_tv = None
if self.cfg.train.loss.pixel_weight > 0.:
loss_pixel = self.pixel_loss_fn(outputs['images_rgb'], render_image)
loss += loss_pixel * self.cfg.train.loss.pixel_weight
if self.cfg.train.loss.perceptual_weight > 0.:
loss_perceptual = self.perceptual_loss_fn(outputs['images_rgb'], render_image)
loss += loss_perceptual * self.cfg.train.loss.perceptual_weight
if self.cfg.train.loss.tv_weight > 0.:
loss_tv = self.tv_loss_fn(outputs['planes'])
loss += loss_tv * self.cfg.train.loss.tv_weight
return outputs, loss, loss_pixel, loss_perceptual, loss_tv
def train_epoch(self, pbar: tqdm, loader: torch.utils.data.DataLoader, profiler: torch.profiler.profile):
self.model.train()
local_step_losses = []
global_step_losses = []
logger.debug(f"======== Starting epoch {self.current_epoch} ========")
for data in loader:
logger.debug(f"======== Starting global step {self.global_step} ========")
with self.accelerator.accumulate(self.model):
# forward to loss
outs, loss, loss_pixel, loss_perceptual, loss_tv = self.forward_loss_local_step(data)
# backward
self.accelerator.backward(loss)
if self.accelerator.sync_gradients and self.cfg.train.optim.clip_grad_norm > 0.:
self.accelerator.clip_grad_norm_(self.model.parameters(), self.cfg.train.optim.clip_grad_norm)
self.optimizer.step()
self.optimizer.zero_grad()
# track local losses
local_step_losses.append(torch.stack([
_loss.detach() if _loss is not None else torch.tensor(float('nan'), device=self.device)
for _loss in [loss, loss_pixel, loss_perceptual, loss_tv]
]))
# track global step
if self.accelerator.sync_gradients:
profiler.step()
self.scheduler.step()
logger.debug(f"======== Scheduler step ========")
self.global_step += 1
global_step_loss = self.accelerator.gather(torch.stack(local_step_losses)).mean(dim=0).cpu()
loss, loss_pixel, loss_perceptual, loss_tv = global_step_loss.unbind()
loss_kwargs = {
'loss': loss.item(),
'loss_pixel': loss_pixel.item(),
'loss_perceptual': loss_perceptual.item(),
'loss_tv': loss_tv.item(),
}
self.log_scalar_kwargs(
step=self.global_step, split='train',
**loss_kwargs
)
self.log_optimizer(step=self.global_step, attrs=['lr'], group_ids=[0, 1])
local_step_losses = []
global_step_losses.append(global_step_loss)
# manage display
pbar.update(1)
description = {
**loss_kwargs,
'lr': self.optimizer.param_groups[0]['lr'],
}
description = '[TRAIN STEP]' + \
', '.join(f'{k}={tqdm.format_num(v)}' for k, v in description.items() if not math.isnan(v))
pbar.set_description(description)
# periodic actions
if self.global_step % self.cfg.saver.checkpoint_global_steps == 0:
self.save_checkpoint()
if self.global_step % self.cfg.val.global_step_period == 0:
self.evaluate()
self.model.train()
if self.global_step % self.cfg.logger.image_monitor.train_global_steps == 0:
self.log_image_monitor(
step=self.global_step, split='train',
renders=outs['images_rgb'].detach()[:self.cfg.logger.image_monitor.samples_per_log].cpu(),
gts=data['render_image'][:self.cfg.logger.image_monitor.samples_per_log].cpu(),
)
# progress control
if self.global_step >= self.N_max_global_steps:
self.accelerator.set_trigger()
break
# track epoch
self.current_epoch += 1
epoch_losses = torch.stack(global_step_losses).mean(dim=0)
epoch_loss, epoch_loss_pixel, epoch_loss_perceptual, epoch_loss_tv = epoch_losses.unbind()
epoch_loss_dict = {
'loss': epoch_loss.item(),
'loss_pixel': epoch_loss_pixel.item(),
'loss_perceptual': epoch_loss_perceptual.item(),
'loss_tv': epoch_loss_tv.item(),
}
self.log_scalar_kwargs(
epoch=self.current_epoch, split='train',
**epoch_loss_dict,
)
logger.info(
f'[TRAIN EPOCH] {self.current_epoch}/{self.cfg.train.epochs}: ' + \
', '.join(f'{k}={tqdm.format_num(v)}' for k, v in epoch_loss_dict.items() if not math.isnan(v))
)
def train(self):
starting_local_step_in_epoch = self.global_step_in_epoch * self.cfg.train.accum_steps
skipped_loader = self.accelerator.skip_first_batches(self.train_loader, starting_local_step_in_epoch)
logger.info(f"======== Skipped {starting_local_step_in_epoch} local batches ========")
with tqdm(
range(0, self.N_max_global_steps),
initial=self.global_step,
disable=(not self.accelerator.is_main_process),
) as pbar:
profiler = torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
schedule=torch.profiler.schedule(
wait=10, warmup=10, active=100,
),
on_trace_ready=torch.profiler.tensorboard_trace_handler(os.path.join(
self.cfg.logger.tracker_root,
self.cfg.experiment.parent, self.cfg.experiment.child,
)),
record_shapes=True,
profile_memory=True,
with_stack=True,
) if self.cfg.logger.enable_profiler else DummyProfiler()
with profiler:
self.optimizer.zero_grad()
for _ in range(self.current_epoch, self.cfg.train.epochs):
loader = skipped_loader or self.train_loader
skipped_loader = None
self.train_epoch(pbar=pbar, loader=loader, profiler=profiler)
if self.accelerator.check_trigger():
break
logger.info(f"======== Training finished at global step {self.global_step} ========")
# final checkpoint and evaluation
self.save_checkpoint()
self.evaluate()
@torch.no_grad()
@torch.compiler.disable
def evaluate(self, epoch: int = None):
self.model.eval()
max_val_batches = self.cfg.val.debug_batches or len(self.val_loader)
running_losses = []
sample_data, sample_outs = None, None
for data in tqdm(self.val_loader, disable=(not self.accelerator.is_main_process), total=max_val_batches):
if len(running_losses) >= max_val_batches:
logger.info(f"======== Early stop validation at {len(running_losses)} batches ========")
break
outs, loss, loss_pixel, loss_perceptual, loss_tv = self.forward_loss_local_step(data)
sample_data, sample_outs = data, outs
running_losses.append(torch.stack([
_loss if _loss is not None else torch.tensor(float('nan'), device=self.device)
for _loss in [loss, loss_pixel, loss_perceptual, loss_tv]
]))
total_losses = self.accelerator.gather(torch.stack(running_losses)).mean(dim=0).cpu()
total_loss, total_loss_pixel, total_loss_perceptual, total_loss_tv = total_losses.unbind()
total_loss_dict = {
'loss': total_loss.item(),
'loss_pixel': total_loss_pixel.item(),
'loss_perceptual': total_loss_perceptual.item(),
'loss_tv': total_loss_tv.item(),
}
if epoch is not None:
self.log_scalar_kwargs(
epoch=epoch, split='val',
**total_loss_dict,
)
logger.info(
f'[VAL EPOCH] {epoch}/{self.cfg.train.epochs}: ' + \
', '.join(f'{k}={tqdm.format_num(v)}' for k, v in total_loss_dict.items() if not math.isnan(v))
)
self.log_image_monitor(
epoch=epoch, split='val',
renders=sample_outs['images_rgb'][:self.cfg.logger.image_monitor.samples_per_log].cpu(),
gts=sample_data['render_image'][:self.cfg.logger.image_monitor.samples_per_log].cpu(),
)
else:
self.log_scalar_kwargs(
step=self.global_step, split='val',
**total_loss_dict,
)
logger.info(
f'[VAL STEP] {self.global_step}/{self.N_max_global_steps}: ' + \
', '.join(f'{k}={tqdm.format_num(v)}' for k, v in total_loss_dict.items() if not math.isnan(v))
)
self.log_image_monitor(
step=self.global_step, split='val',
renders=sample_outs['images_rgb'][:self.cfg.logger.image_monitor.samples_per_log].cpu(),
gts=sample_data['render_image'][:self.cfg.logger.image_monitor.samples_per_log].cpu(),
)
@Trainer.control('on_main_process')
def log_image_monitor(
self, epoch: int = None, step: int = None, split: str = None,
renders: torch.Tensor = None, gts: torch.Tensor = None,
):
M = renders.shape[1]
merged = torch.stack([renders, gts], dim=1)[0].view(-1, *renders.shape[2:])
renders, gts = renders.view(-1, *renders.shape[2:]), gts.view(-1, *gts.shape[2:])
renders, gts, merged = make_grid(renders, nrow=M), make_grid(gts, nrow=M), make_grid(merged, nrow=M)
log_type, log_progress = self._get_str_progress(epoch, step)
split = f'/{split}' if split else ''
self.log_images({
f'Images_split{split}/rendered': renders.unsqueeze(0),
f'Images_split{split}/gt': gts.unsqueeze(0),
f'Images_merged{split}': merged.unsqueeze(0),
}, log_progress)
# 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 os
import time
import math
import argparse
import shutil
import torch
import safetensors
from omegaconf import OmegaConf
from abc import abstractmethod
from contextlib import contextmanager
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed
from openlrm.utils.logging import configure_logger
from openlrm.utils.compile import configure_dynamo
from openlrm.runners.abstract import Runner
logger = get_logger(__name__)
def parse_configs():
# Define argparse arguments
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default='./assets/config.yaml')
args, unknown = parser.parse_known_args()
# Load configuration file
cfg = OmegaConf.load(args.config)
# Override with command-line arguments
cli_cfg = OmegaConf.from_cli(unknown)
cfg = OmegaConf.merge(cfg, cli_cfg)
return cfg
class Trainer(Runner):
def __init__(self):
super().__init__()
self.cfg = parse_configs()
self.timestamp = time.strftime("%Y%m%d-%H%M%S")
self.accelerator = Accelerator(
mixed_precision=self.cfg.train.mixed_precision,
gradient_accumulation_steps=self.cfg.train.accum_steps,
log_with=tuple(self.cfg.logger.trackers),
project_config=ProjectConfiguration(
logging_dir=self.cfg.logger.tracker_root,
),
use_seedable_sampler=True,
kwargs_handlers=[
DistributedDataParallelKwargs(
find_unused_parameters=self.cfg.train.find_unused_parameters,
),
],
)
set_seed(self.cfg.experiment.seed, device_specific=True)
with self.accelerator.main_process_first():
configure_logger(
stream_level=self.cfg.logger.stream_level,
log_level=self.cfg.logger.log_level,
file_path=os.path.join(
self.cfg.logger.log_root,
self.cfg.experiment.parent, self.cfg.experiment.child,
f"{self.timestamp}.log",
) if self.accelerator.is_main_process else None,
)
logger.info(self.accelerator.state, main_process_only=False, in_order=True)
configure_dynamo(dict(self.cfg.compile))
# attributes with defaults
self.model : torch.nn.Module = None
self.optimizer: torch.optim.Optimizer = None
self.scheduler: torch.optim.lr_scheduler.LRScheduler = None
self.train_loader: torch.utils.data.DataLoader = None
self.val_loader: torch.utils.data.DataLoader = None
self.N_max_global_steps: int = None
self.N_global_steps_per_epoch: int = None
self.global_step: int = 0
self.current_epoch: int = 0
def __enter__(self):
self.accelerator.init_trackers(
project_name=f"{self.cfg.experiment.parent}/{self.cfg.experiment.child}",
)
self.prepare_everything()
self.log_inital_info()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.accelerator.end_training()
@staticmethod
def control(option: str = None, synchronized: bool = False):
def decorator(func):
def wrapper(self, *args, **kwargs):
if option is None or hasattr(self.accelerator, option):
accelerated_func = getattr(self.accelerator, option)(func) if option is not None else func
result = accelerated_func(self, *args, **kwargs)
if synchronized:
self.accelerator.wait_for_everyone()
return result
else:
raise AttributeError(f"Accelerator has no attribute {option}")
return wrapper
return decorator
@contextmanager
def exec_in_order(self):
for rank in range(self.accelerator.num_processes):
try:
if self.accelerator.process_index == rank:
yield
finally:
self.accelerator.wait_for_everyone()
@property
def device(self):
return self.accelerator.device
@property
def is_distributed(self) -> bool:
return self.accelerator.num_processes > 1
def prepare_everything(self, is_dist_validation: bool = True):
# prepare with accelerator
if is_dist_validation:
self.model, self.optimizer, self.train_loader, self.val_loader = \
self.accelerator.prepare(
self.model, self.optimizer, self.train_loader, self.val_loader,
)
else:
self.model, self.optimizer, self.train_loader = \
self.accelerator.prepare(
self.model, self.optimizer, self.train_loader,
)
self.accelerator.register_for_checkpointing(self.scheduler)
# prepare stats
N_total_batch_size = self.cfg.train.batch_size * self.accelerator.num_processes * self.cfg.train.accum_steps
self.N_global_steps_per_epoch = math.ceil(len(self.train_loader) / self.cfg.train.accum_steps)
self.N_max_global_steps = self.N_global_steps_per_epoch * self.cfg.train.epochs
if self.cfg.train.debug_global_steps is not None:
logger.warning(f"Overriding max global steps from {self.N_max_global_steps} to {self.cfg.train.debug_global_steps}")
self.N_max_global_steps = self.cfg.train.debug_global_steps
logger.info(f"======== Statistics ========")
logger.info(f"** N_max_global_steps: {self.N_max_global_steps}")
logger.info(f"** N_total_batch_size: {N_total_batch_size}")
logger.info(f"** N_epochs: {self.cfg.train.epochs}")
logger.info(f"** N_global_steps_per_epoch: {self.N_global_steps_per_epoch}")
logger.debug(f"** Prepared loader length: {len(self.train_loader)}")
logger.info(f"** Distributed validation: {is_dist_validation}")
logger.info(f"============================")
logger.info(f"======== Trainable parameters ========")
logger.info(f"** Total: {sum(p.numel() for p in self.model.parameters() if p.requires_grad)}")
for sub_name, sub_module in self.accelerator.unwrap_model(self.model).named_children():
logger.info(f"** {sub_name}: {sum(p.numel() for p in sub_module.parameters() if p.requires_grad)}")
logger.info(f"=====================================")
self.accelerator.wait_for_everyone()
# load checkpoint or model
self.load_ckpt_or_auto_resume_(self.cfg)
# register hooks
self.register_hooks()
@abstractmethod
def register_hooks(self):
pass
def auto_resume_(self, cfg) -> bool:
ckpt_root = os.path.join(
cfg.saver.checkpoint_root,
cfg.experiment.parent, cfg.experiment.child,
)
if not os.path.exists(ckpt_root):
return False
ckpt_dirs = os.listdir(ckpt_root)
if len(ckpt_dirs) == 0:
return False
ckpt_dirs.sort()
latest_ckpt = ckpt_dirs[-1]
latest_ckpt_dir = os.path.join(ckpt_root, latest_ckpt)
logger.info(f"======== Auto-resume from {latest_ckpt_dir} ========")
self.accelerator.load_state(latest_ckpt_dir)
self.global_step = int(latest_ckpt)
self.current_epoch = self.global_step // self.N_global_steps_per_epoch
return True
def load_model_(self, cfg):
logger.info(f"======== Loading model from {cfg.saver.load_model} ========")
safetensors.torch.load_model(
self.accelerator.unwrap_model(self.model),
cfg.saver.load_model,
strict=True,
)
logger.info(f"======== Model loaded ========")
@control(synchronized=True)
def load_ckpt_or_auto_resume_(self, cfg):
# auto resume has higher priority, load model from path if auto resume is not available
# cfg.saver.auto_resume and cfg.saver.load_model
if cfg.saver.auto_resume:
successful_resume = self.auto_resume_(cfg)
if successful_resume:
return
if cfg.saver.load_model:
successful_load = self.load_model_(cfg)
if successful_load:
return
logger.debug(f"======== No checkpoint or model is loaded ========")
@control('on_main_process', synchronized=True)
def save_checkpoint(self):
ckpt_dir = os.path.join(
self.cfg.saver.checkpoint_root,
self.cfg.experiment.parent, self.cfg.experiment.child,
f"{self.global_step:06d}",
)
self.accelerator.save_state(output_dir=ckpt_dir, safe_serialization=True)
logger.info(f"======== Saved checkpoint at global step {self.global_step} ========")
# manage stratified checkpoints
ckpt_dirs = os.listdir(os.path.dirname(ckpt_dir))
ckpt_dirs.sort()
max_ckpt = int(ckpt_dirs[-1])
ckpt_base = int(self.cfg.saver.checkpoint_keep_level)
ckpt_period = self.cfg.saver.checkpoint_global_steps
logger.debug(f"Checkpoint base: {ckpt_base}")
logger.debug(f"Checkpoint period: {ckpt_period}")
cur_order = ckpt_base ** math.floor(math.log(max_ckpt // ckpt_period, ckpt_base))
cur_idx = 0
while cur_order > 0:
cur_digit = max_ckpt // ckpt_period // cur_order % ckpt_base
while cur_idx < len(ckpt_dirs) and int(ckpt_dirs[cur_idx]) // ckpt_period // cur_order % ckpt_base < cur_digit:
if int(ckpt_dirs[cur_idx]) // ckpt_period % cur_order != 0:
shutil.rmtree(os.path.join(os.path.dirname(ckpt_dir), ckpt_dirs[cur_idx]))
logger.info(f"Removed checkpoint {ckpt_dirs[cur_idx]}")
cur_idx += 1
cur_order //= ckpt_base
@property
def global_step_in_epoch(self):
return self.global_step % self.N_global_steps_per_epoch
@abstractmethod
def _build_model(self):
pass
@abstractmethod
def _build_optimizer(self):
pass
@abstractmethod
def _build_scheduler(self):
pass
@abstractmethod
def _build_dataloader(self):
pass
@abstractmethod
def _build_loss_fn(self):
pass
@abstractmethod
def train(self):
pass
@abstractmethod
def evaluate(self):
pass
@staticmethod
def _get_str_progress(epoch: int = None, step: int = None):
if epoch is not None:
log_type = 'epoch'
log_progress = epoch
elif step is not None:
log_type = 'step'
log_progress = step
else:
raise ValueError('Either epoch or step must be provided')
return log_type, log_progress
@control('on_main_process')
def log_scalar_kwargs(self, epoch: int = None, step: int = None, split: str = None, **scalar_kwargs):
log_type, log_progress = self._get_str_progress(epoch, step)
split = f'/{split}' if split else ''
for key, value in scalar_kwargs.items():
self.accelerator.log({f'{key}{split}/{log_type}': value}, log_progress)
@control('on_main_process')
def log_images(self, values: dict, step: int | None = None, log_kwargs: dict | None = {}):
for tracker in self.accelerator.trackers:
if hasattr(tracker, 'log_images'):
tracker.log_images(values, step=step, **log_kwargs.get(tracker.name, {}))
@control('on_main_process')
def log_optimizer(self, epoch: int = None, step: int = None, attrs: list[str] = [], group_ids: list[int] = []):
log_type, log_progress = self._get_str_progress(epoch, step)
assert self.optimizer is not None, 'Optimizer is not initialized'
if not attrs:
logger.warning('No optimizer attributes are provided, nothing will be logged')
if not group_ids:
logger.warning('No optimizer group ids are provided, nothing will be logged')
for attr in attrs:
assert attr in ['lr', 'momentum', 'weight_decay'], f'Invalid optimizer attribute {attr}'
for group_id in group_ids:
self.accelerator.log({f'opt/{attr}/{group_id}': self.optimizer.param_groups[group_id][attr]}, log_progress)
@control('on_main_process')
def log_inital_info(self):
assert self.model is not None, 'Model is not initialized'
assert self.optimizer is not None, 'Optimizer is not initialized'
assert self.scheduler is not None, 'Scheduler is not initialized'
self.accelerator.log({'Config': "```\n" + OmegaConf.to_yaml(self.cfg) + "\nCode · "})
self.accelerator.log({'Model': "\n" + str(self.model) + "\n
Code · "})
self.accelerator.log({'Optimizer': "\n" + str(self.optimizer) + "\n
Code · "})
self.accelerator.log({'Scheduler': "\n" + str(self.scheduler) + "\n```"})
def run(self):
self.train()
# 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.
from .lrm import LRMTrainer
# 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
from abc import abstractmethod
from accelerate import Accelerator
from accelerate.logging import get_logger
from openlrm.runners.abstract import Runner
logger = get_logger(__name__)
class Inferrer(Runner):
EXP_TYPE: str = None
def __init__(self):
super().__init__()
torch._dynamo.config.disable = True
self.accelerator = Accelerator()
self.model : torch.nn.Module = None
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
pass
@property
def device(self):
return self.accelerator.device
@abstractmethod
def _build_model(self, cfg):
pass
@abstractmethod
def infer_single(self, *args, **kwargs):
pass
@abstractmethod
def infer(self):
pass
def run(self):
self.infer()
# 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 numpy as np
from PIL import Image
from omegaconf import OmegaConf
from tqdm.auto import tqdm
from accelerate.logging import get_logger
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 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)
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.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('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.logger,
log_level=self.cfg.logger,
)
self.model = self._build_model(self.cfg).to(self.device)
def _build_model(self, cfg):
from openlrm.models import model_dict
hf_model_cls = wrap_model_hub(model_dict[self.EXP_TYPE])
model = hf_model_cls.from_pretrained(cfg.model_name)
return model
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):
N = image.shape[0]
source_camera = self._default_source_camera(dist_to_center=source_cam_dist, batch_size=N, device=self.device)
planes = self.model.forward_planes(image, source_camera)
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):
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':
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):
source_size = self.cfg.source_size
render_size = self.cfg.render_size
render_views = self.cfg.render_views
render_fps = self.cfg.render_fps
mesh_size = self.cfg.mesh_size
mesh_thres = self.cfg.mesh_thres
frame_size = self.cfg.frame_size
source_cam_dist = self.cfg.source_cam_dist if source_cam_dist is None else source_cam_dist
# 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)
with torch.no_grad():
planes = self.infer_planes(image, source_cam_dist=source_cam_dist)
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)
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 infer(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
for root, dirs, files in os.walk(self.cfg.image_input):
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]
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,
)
# 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.
from .lrm import LRMInferrer
# 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 torch.nn as nn
__all__ = ['TVLoss']
class TVLoss(nn.Module):
"""
Total variance loss.
"""
def __init__(self):
super().__init__()
def numel_excluding_first_dim(self, x):
return x.numel() // x.shape[0]
@torch.compile
def forward(self, x):
"""
Assume batched and channel first with inner sizes.
Args:
x: [N, M, C, H, W]
Returns:
Mean-reduced TV loss with element-level scaling.
"""
N, M, C, H, W = x.shape
x = x.reshape(N*M, C, H, W)
diff_i = x[..., 1:, :] - x[..., :-1, :]
diff_j = x[..., :, 1:] - x[..., :, :-1]
div_i = self.numel_excluding_first_dim(diff_i)
div_j = self.numel_excluding_first_dim(diff_j)
tv_i = diff_i.pow(2).sum(dim=[1,2,3]) / div_i
tv_j = diff_j.pow(2).sum(dim=[1,2,3]) / div_j
tv = tv_i + tv_j
batch_tv = tv.reshape(N, M).mean(dim=1)
all_tv = batch_tv.mean()
return all_tv
# 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 torch.nn as nn
__all__ = ['LPIPSLoss']
class LPIPSLoss(nn.Module):
"""
Compute LPIPS loss between two images.
"""
def __init__(self, device, prefech: bool = False):
super().__init__()
self.device = device
self.cached_models = {}
if prefech:
self.prefetch_models()
def _get_model(self, model_name: str):
if model_name not in self.cached_models:
import warnings
with warnings.catch_warnings():
warnings.filterwarnings('ignore', category=UserWarning)
import lpips
_model = lpips.LPIPS(net=model_name, eval_mode=True, verbose=False).to(self.device)
_model = torch.compile(_model)
self.cached_models[model_name] = _model
return self.cached_models[model_name]
def prefetch_models(self):
_model_names = ['alex', 'vgg']
for model_name in _model_names:
self._get_model(model_name)
def forward(self, x, y, is_training: bool = True):
"""
Assume images are 0-1 scaled and channel first.
Args:
x: [N, M, C, H, W]
y: [N, M, C, H, W]
is_training: whether to use VGG or AlexNet.
Returns:
Mean-reduced LPIPS loss across batch.
"""
model_name = 'vgg' if is_training else 'alex'
loss_fn = self._get_model(model_name)
N, M, C, H, W = x.shape
x = x.reshape(N*M, C, H, W)
y = y.reshape(N*M, C, H, W)
image_loss = loss_fn(x, y, normalize=True).mean(dim=[1, 2, 3])
batch_loss = image_loss.reshape(N, M).mean(dim=1)
all_loss = batch_loss.mean()
return all_loss
# 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 torch.nn as nn
__all__ = ['PixelLoss']
class PixelLoss(nn.Module):
"""
Pixel-wise loss between two images.
"""
def __init__(self, option: str = 'mse'):
super().__init__()
self.loss_fn = self._build_from_option(option)
@staticmethod
def _build_from_option(option: str, reduction: str = 'none'):
if option == 'mse':
return nn.MSELoss(reduction=reduction)
elif option == 'l1':
return nn.L1Loss(reduction=reduction)
else:
raise NotImplementedError(f'Unknown pixel loss option: {option}')
@torch.compile
def forward(self, x, y):
"""
Assume images are channel first.
Args:
x: [N, M, C, H, W]
y: [N, M, C, H, W]
Returns:
Mean-reduced pixel loss across batch.
"""
N, M, C, H, W = x.shape
x = x.reshape(N*M, C, H, W)
y = y.reshape(N*M, C, H, W)
image_loss = self.loss_fn(x, y).mean(dim=[1, 2, 3])
batch_loss = image_loss.reshape(N, M).mean(dim=1)
all_loss = batch_loss.mean()
return all_loss
# 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.
from .pixelwise import *
from .perceptual import *
from .tvloss import *
# 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 torch.nn as nn
from accelerate.logging import get_logger
from .embedder import CameraEmbedder
from .transformer import TransformerDecoder
from .rendering.synthesizer import TriplaneSynthesizer
logger = get_logger(__name__)
class ModelLRM(nn.Module):
"""
Full model of the basic single-view 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_type: str = 'dino',
encoder_model_name: str = 'facebook/dino-vitb16', encoder_feat_dim: int = 768):
super().__init__()
# attributes
self.encoder_feat_dim = encoder_feat_dim
self.camera_embed_dim = camera_embed_dim
self.triplane_low_res = triplane_low_res
self.triplane_high_res = triplane_high_res
self.triplane_dim = triplane_dim
# modules
self.encoder = self._encoder_fn(encoder_type)(
model_name=encoder_model_name,
freeze=encoder_freeze,
)
self.camera_embedder = CameraEmbedder(
raw_dim=12+4, embed_dim=camera_embed_dim,
)
# initialize pos_embed with 1/sqrt(dim) * N(0, 1)
self.pos_embed = nn.Parameter(torch.randn(1, 3*triplane_low_res**2, transformer_dim) * (1. / transformer_dim) ** 0.5)
self.transformer = TransformerDecoder(
block_type='cond_mod',
num_layers=transformer_layers, num_heads=transformer_heads,
inner_dim=transformer_dim, cond_dim=encoder_feat_dim, mod_dim=camera_embed_dim,
)
self.upsampler = nn.ConvTranspose2d(transformer_dim, triplane_dim, kernel_size=2, stride=2, padding=0)
self.synthesizer = TriplaneSynthesizer(
triplane_dim=triplane_dim, samples_per_ray=rendering_samples_per_ray,
)
@staticmethod
def _encoder_fn(encoder_type: str):
encoder_type = encoder_type.lower()
assert encoder_type in ['dino', 'dinov2'], "Unsupported encoder type"
if encoder_type == 'dino':
from .encoders.dino_wrapper import DinoWrapper
logger.info("Using DINO as the encoder")
return DinoWrapper
elif encoder_type == 'dinov2':
from .encoders.dinov2_wrapper import Dinov2Wrapper
logger.info("Using DINOv2 as the encoder")
return Dinov2Wrapper
def forward_transformer(self, image_feats, camera_embeddings):
assert image_feats.shape[0] == camera_embeddings.shape[0], \
"Batch size mismatch for image_feats and camera_embeddings!"
N = image_feats.shape[0]
x = self.pos_embed.repeat(N, 1, 1) # [N, L, D]
x = self.transformer(
x,
cond=image_feats,
mod=camera_embeddings,
)
return x
def reshape_upsample(self, tokens):
N = tokens.shape[0]
H = W = self.triplane_low_res
x = tokens.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.upsampler(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()
return x
@torch.compile
def forward_planes(self, image, camera):
# image: [N, C_img, H_img, W_img]
# camera: [N, D_cam_raw]
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
tokens = self.forward_transformer(image_feats, camera_embeddings)
planes = self.reshape_upsample(tokens)
assert planes.shape[0] == N, "Batch size mismatch for planes"
assert planes.shape[1] == 3, "Planes should have 3 channels"
return planes
def forward(self, image, source_camera, render_cameras, render_anchors, render_resolutions, render_bg_colors, render_region_size: int):
# image: [N, C_img, H_img, W_img]
# source_camera: [N, D_cam_raw]
# render_cameras: [N, M, D_cam_render]
# render_anchors: [N, M, 2]
# render_resolutions: [N, M, 1]
# render_bg_colors: [N, M, 1]
# render_region_size: int
assert image.shape[0] == source_camera.shape[0], "Batch size mismatch for image and source_camera"
assert image.shape[0] == render_cameras.shape[0], "Batch size mismatch for image and render_cameras"
assert image.shape[0] == render_anchors.shape[0], "Batch size mismatch for image and render_anchors"
assert image.shape[0] == render_bg_colors.shape[0], "Batch size mismatch for image and render_bg_colors"
N, M = render_cameras.shape[:2]
planes = self.forward_planes(image, source_camera)
# render target views
render_results = self.synthesizer(planes, render_cameras, render_anchors, render_resolutions, render_bg_colors, render_region_size)
assert render_results['images_rgb'].shape[0] == N, "Batch size mismatch for render_results"
assert render_results['images_rgb'].shape[1] == M, "Number of rendered views should be consistent with render_cameras"
return {
'planes': planes,
**render_results,
}
# 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 torch.nn as nn
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),
)
@torch.compile
def forward(self, x):
return self.mlp(x)
# 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.
from functools import partial
import torch
import torch.nn as nn
from accelerate.logging import get_logger
logger = get_logger(__name__)
class TransformerDecoder(nn.Module):
"""
Transformer blocks that process the input and optionally use condition and modulation.
"""
def __init__(self, block_type: str,
num_layers: int, num_heads: int,
inner_dim: int, cond_dim: int = None, mod_dim: int = None,
eps: float = 1e-6):
super().__init__()
self.block_type = block_type
self.layers = nn.ModuleList([
self._block_fn(inner_dim, cond_dim, mod_dim)(
num_heads=num_heads,
eps=eps,
)
for _ in range(num_layers)
])
self.norm = nn.LayerNorm(inner_dim, eps=eps)
@property
def block_type(self):
return self._block_type
@block_type.setter
def block_type(self, block_type):
assert block_type in ['basic', 'cond', 'mod', 'cond_mod'], \
f"Unsupported block type: {block_type}"
self._block_type = block_type
def _block_fn(self, inner_dim, cond_dim, mod_dim):
assert inner_dim is not None, f"inner_dim must always be specified"
if self.block_type == 'basic':
assert cond_dim is None and mod_dim is None, \
f"Condition and modulation are not supported for BasicBlock"
from .block import BasicBlock
logger.debug(f"Using BasicBlock")
return partial(BasicBlock, inner_dim=inner_dim)
elif self.block_type == 'cond':
assert cond_dim is not None, f"Condition dimension must be specified for ConditionBlock"
assert mod_dim is None, f"Modulation dimension is not supported for ConditionBlock"
from .block import ConditionBlock
logger.debug(f"Using ConditionBlock")
return partial(ConditionBlock, inner_dim=inner_dim, cond_dim=cond_dim)
elif self.block_type == 'mod':
logger.error(f"modulation without condition is not implemented")
raise NotImplementedError(f"modulation without condition is not implemented")
elif self.block_type == 'cond_mod':
assert cond_dim is not None and mod_dim is not None, \
f"Condition and modulation dimensions must be specified for ConditionModulationBlock"
from .block import ConditionModulationBlock
logger.debug(f"Using ConditionModulationBlock")
return partial(ConditionModulationBlock, inner_dim=inner_dim, cond_dim=cond_dim, mod_dim=mod_dim)
else:
raise ValueError(f"Unsupported block type during runtime: {self.block_type}")
def assert_runtime_integrity(self, x: torch.Tensor, cond: torch.Tensor, mod: torch.Tensor):
assert x is not None, f"Input tensor must be specified"
if self.block_type == 'basic':
assert cond is None and mod is None, \
f"Condition and modulation are not supported for BasicBlock"
elif self.block_type == 'cond':
assert cond is not None and mod is None, \
f"Condition must be specified and modulation is not supported for ConditionBlock"
elif self.block_type == 'mod':
raise NotImplementedError(f"modulation without condition is not implemented")
else:
assert cond is not None and mod is not None, \
f"Condition and modulation must be specified for ConditionModulationBlock"
def forward_layer(self, layer: nn.Module, x: torch.Tensor, cond: torch.Tensor, mod: torch.Tensor):
if self.block_type == 'basic':
return layer(x)
elif self.block_type == 'cond':
return layer(x, cond)
elif self.block_type == 'mod':
return layer(x, mod)
else:
return layer(x, cond, mod)
def forward(self, x: torch.Tensor, cond: torch.Tensor = None, mod: torch.Tensor = None):
# x: [N, L, D]
# cond: [N, L_cond, D_cond] or None
# mod: [N, D_mod] or None
self.assert_runtime_integrity(x, cond, mod)
for layer in self.layers:
x = self.forward_layer(layer, x, cond, mod)
x = self.norm(x)
return x
# 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.nn as nn
from .modulate import ModLN
class BasicBlock(nn.Module):
"""
Transformer block that is in its simplest form.
Designed for PF-LRM architecture.
"""
# Block contains a self-attention layer and an MLP
def __init__(self, inner_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 = nn.LayerNorm(inner_dim, eps=eps)
self.self_attn = nn.MultiheadAttention(
embed_dim=inner_dim, num_heads=num_heads,
dropout=attn_drop, bias=attn_bias, batch_first=True)
self.norm2 = nn.LayerNorm(inner_dim, eps=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):
# x: [N, L, D]
before_sa = self.norm1(x)
x = x + self.self_attn(before_sa, before_sa, before_sa, need_weights=False)[0]
x = x + self.mlp(self.norm2(x))
return x
class ConditionBlock(nn.Module):
"""
Transformer block that takes in a cross-attention condition.
Designed for SparseLRM architecture.
"""
# Block contains a cross-attention layer, a self-attention layer, and an MLP
def __init__(self, inner_dim: int, cond_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 = nn.LayerNorm(inner_dim, eps=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 = nn.LayerNorm(inner_dim, eps=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 = nn.LayerNorm(inner_dim, eps=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):
# x: [N, L, D]
# cond: [N, L_cond, D_cond]
x = x + self.cross_attn(self.norm1(x), cond, cond, need_weights=False)[0]
before_sa = self.norm2(x)
x = x + self.self_attn(before_sa, before_sa, before_sa, need_weights=False)[0]
x = x + self.mlp(self.norm3(x))
return x
class ConditionModulationBlock(nn.Module):
"""
Transformer block that takes in a cross-attention condition and another modulation vector applied to sub-blocks.
Designed for raw LRM architecture.
"""
# Block contains a cross-attention layer, a self-attention layer, and an 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
# 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 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: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:
shift, scale = self.mlp(mod).chunk(2, dim=-1) # [N, D]
return self.modulate(self.norm(x), shift, scale) # [N, L, D]
# 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.
from .modeling_lrm import ModelLRM
model_dict = {
'lrm': ModelLRM,
}
# 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 torch.nn as nn
from accelerate.logging import get_logger
logger = get_logger(__name__)
class Dinov2Wrapper(nn.Module):
"""
Dino v2 wrapper using original implementation, hacked with modulation.
"""
def __init__(self, model_name: str, modulation_dim: int = None, freeze: bool = True):
super().__init__()
self.modulation_dim = modulation_dim
self.model = self._build_dinov2(model_name, modulation_dim=modulation_dim)
if freeze:
if modulation_dim is not None:
raise ValueError("Modulated Dinov2 requires training, freezing is not allowed.")
self._freeze()
def _freeze(self):
logger.warning(f"======== Freezing Dinov2Wrapper ========")
self.model.eval()
for name, param in self.model.named_parameters():
param.requires_grad = False
@staticmethod
def _build_dinov2(model_name: str, modulation_dim: int = None, pretrained: bool = True):
from importlib import import_module
dinov2_hub = import_module(".dinov2.hub.backbones", package=__package__)
model_fn = getattr(dinov2_hub, model_name)
logger.debug(f"Modulation dim for Dinov2 is {modulation_dim}.")
model = model_fn(modulation_dim=modulation_dim, pretrained=pretrained)
return model
@torch.compile
def forward(self, image: torch.Tensor, mod: torch.Tensor = None):
# image: [N, C, H, W]
# mod: [N, D] or None
# RGB image with [0,1] scale and properly sized
if self.modulation_dim is None:
assert mod is None, "Unexpected modulation input in dinov2 forward."
outs = self.model(image, is_training=True)
else:
assert mod is not None, "Modulation input is required in modulated dinov2 forward."
outs = self.model(image, mod=mod, is_training=True)
ret = torch.cat([
outs["x_norm_clstoken"].unsqueeze(dim=1),
outs["x_norm_patchtokens"],
], dim=1)
return ret
# 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.
#
# Empty
# 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 torch.nn as nn
from transformers import ViTImageProcessor, ViTModel
from accelerate.logging import get_logger
logger = get_logger(__name__)
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()
@torch.compile
def forward_model(self, inputs):
return self.model(**inputs, interpolate_pos_encoding=True)
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, return_tensors="pt", do_rescale=False, do_resize=False).to(self.model.device)
# This resampling of positional embedding uses bicubic interpolation
outputs = self.forward_model(inputs)
last_hidden_states = outputs.last_hidden_state
return last_hidden_states
def _freeze(self):
logger.warning(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:
model = ViTModel.from_pretrained(model_name, add_pooling_layer=False)
processor = ViTImageProcessor.from_pretrained(model_name)
return model, processor
except requests.exceptions.ProxyError as err:
if proxy_error_retries > 0:
print(f"Huggingface ProxyError: Retrying ({proxy_error_retries}) 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
# 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.
#
# Empty
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
# References:
# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
# ******************************************************************************
# Code modified by Zexin He in 2023-2024.
# Modifications are marked with clearly visible comments
# licensed under the Apache License, Version 2.0.
# ******************************************************************************
from functools import partial
import math
import logging
from typing import Sequence, Tuple, Union, Callable
import torch
import torch.nn as nn
import torch.utils.checkpoint
from torch.nn.init import trunc_normal_
# ********** Modified by Zexin He in 2023-2024 **********
# Avoid using nested tensor for now, deprecating usage of NestedTensorBlock
from ..layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, Block, BlockWithModulation
# ********************************************************
logger = logging.getLogger("dinov2")
def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
if not depth_first and include_root:
fn(module=module, name=name)
for child_name, child_module in module.named_children():
child_name = ".".join((name, child_name)) if name else child_name
named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
if depth_first and include_root:
fn(module=module, name=name)
return module
class BlockChunk(nn.ModuleList):
def forward(self, x):
for b in self:
x = b(x)
return x
class DinoVisionTransformer(nn.Module):
def __init__(
self,
img_size=224,
patch_size=16,
in_chans=3,
embed_dim=768,
depth=12,
num_heads=12,
mlp_ratio=4.0,
qkv_bias=True,
ffn_bias=True,
proj_bias=True,
drop_path_rate=0.0,
drop_path_uniform=False,
init_values=None, # for layerscale: None or 0 => no layerscale
embed_layer=PatchEmbed,
act_layer=nn.GELU,
block_fn=Block,
# ********** Modified by Zexin He in 2023-2024 **********
modulation_dim: int = None,
# ********************************************************
ffn_layer="mlp",
block_chunks=1,
num_register_tokens=0,
interpolate_antialias=False,
interpolate_offset=0.1,
):
"""
Args:
img_size (int, tuple): input image size
patch_size (int, tuple): patch size
in_chans (int): number of input channels
embed_dim (int): embedding dimension
depth (int): depth of transformer
num_heads (int): number of attention heads
mlp_ratio (int): ratio of mlp hidden dim to embedding dim
qkv_bias (bool): enable bias for qkv if True
proj_bias (bool): enable bias for proj in attn if True
ffn_bias (bool): enable bias for ffn if True
drop_path_rate (float): stochastic depth rate
drop_path_uniform (bool): apply uniform drop rate across blocks
weight_init (str): weight init scheme
init_values (float): layer-scale init values
embed_layer (nn.Module): patch embedding layer
act_layer (nn.Module): MLP activation layer
block_fn (nn.Module): transformer block class
ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
num_register_tokens: (int) number of extra cls tokens (so-called "registers")
interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
"""
super().__init__()
# ********** Modified by Zexin He in 2023-2024 **********
block_norm_layer = None
if modulation_dim is not None:
from ....modulate import ModLN
block_norm_layer = partial(ModLN, mod_dim=modulation_dim)
else:
block_norm_layer = nn.LayerNorm
block_norm_layer = partial(block_norm_layer, eps=1e-6)
# ********************************************************
norm_layer = partial(nn.LayerNorm, eps=1e-6)
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
self.num_tokens = 1
self.n_blocks = depth
self.num_heads = num_heads
self.patch_size = patch_size
self.num_register_tokens = num_register_tokens
self.interpolate_antialias = interpolate_antialias
self.interpolate_offset = interpolate_offset
self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
num_patches = self.patch_embed.num_patches
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
assert num_register_tokens >= 0
self.register_tokens = (
nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
)
if drop_path_uniform is True:
dpr = [drop_path_rate] * depth
else:
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
if ffn_layer == "mlp":
logger.info("using MLP layer as FFN")
ffn_layer = Mlp
elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
logger.info("using SwiGLU layer as FFN")
ffn_layer = SwiGLUFFNFused
elif ffn_layer == "identity":
logger.info("using Identity layer as FFN")
def f(*args, **kwargs):
return nn.Identity()
ffn_layer = f
else:
raise NotImplementedError
blocks_list = [
block_fn(
dim=embed_dim,
num_heads=num_heads,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
proj_bias=proj_bias,
ffn_bias=ffn_bias,
drop_path=dpr[i],
# ********** Modified by Zexin He in 2023-2024 **********
norm_layer=block_norm_layer,
# ********************************************************
act_layer=act_layer,
ffn_layer=ffn_layer,
init_values=init_values,
)
for i in range(depth)
]
if block_chunks > 0:
self.chunked_blocks = True
chunked_blocks = []
chunksize = depth // block_chunks
for i in range(0, depth, chunksize):
# this is to keep the block index consistent if we chunk the block list
chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
else:
self.chunked_blocks = False
self.blocks = nn.ModuleList(blocks_list)
self.norm = norm_layer(embed_dim)
self.head = nn.Identity()
# ********** Modified by Zexin He in 2023-2024 **********
# hacking unused mask_token for better DDP
# self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
# ********************************************************
self.init_weights()
def init_weights(self):
trunc_normal_(self.pos_embed, std=0.02)
nn.init.normal_(self.cls_token, std=1e-6)
if self.register_tokens is not None:
nn.init.normal_(self.register_tokens, std=1e-6)
named_apply(init_weights_vit_timm, self)
def interpolate_pos_encoding(self, x, w, h):
previous_dtype = x.dtype
npatch = x.shape[1] - 1
N = self.pos_embed.shape[1] - 1
if npatch == N and w == h:
return self.pos_embed
pos_embed = self.pos_embed.float()
class_pos_embed = pos_embed[:, 0]
patch_pos_embed = pos_embed[:, 1:]
dim = x.shape[-1]
w0 = w // self.patch_size
h0 = h // self.patch_size
# we add a small number to avoid floating point error in the interpolation
# see discussion at https://github.com/facebookresearch/dino/issues/8
w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
sqrt_N = math.sqrt(N)
sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
patch_pos_embed = nn.functional.interpolate(
patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
scale_factor=(sx, sy),
mode="bicubic",
antialias=self.interpolate_antialias,
)
assert int(w0) == patch_pos_embed.shape[-2]
assert int(h0) == patch_pos_embed.shape[-1]
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
def prepare_tokens_with_masks(self, x, masks=None):
B, nc, w, h = x.shape
x = self.patch_embed(x)
if masks is not None:
# ********** Modified by Zexin He in 2023-2024 **********
raise NotImplementedError("Masking is not supported in hacked DINOv2")
# x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
# ********************************************************
x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
x = x + self.interpolate_pos_encoding(x, w, h)
if self.register_tokens is not None:
x = torch.cat(
(
x[:, :1],
self.register_tokens.expand(x.shape[0], -1, -1),
x[:, 1:],
),
dim=1,
)
return x
def forward_features_list(self, x_list, masks_list):
x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
for blk in self.blocks:
x = blk(x)
all_x = x
output = []
for x, masks in zip(all_x, masks_list):
x_norm = self.norm(x)
output.append(
{
"x_norm_clstoken": x_norm[:, 0],
"x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
"x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
"x_prenorm": x,
"masks": masks,
}
)
return output
# ********** Modified by Zexin He in 2023-2024 **********
def forward_features(self, x, masks=None, mod=None):
if isinstance(x, list):
raise DeprecationWarning("forward_features_list is deprecated, use forward_features")
return self.forward_features_list(x, masks)
x = self.prepare_tokens_with_masks(x, masks)
if mod is None:
for blk in self.blocks:
x = blk(x)
else:
for blk in self.blocks:
x = blk(x, mod)
x_norm = self.norm(x)
return {
"x_norm_clstoken": x_norm[:, 0],
"x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
"x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
"x_prenorm": x,
"masks": masks,
}
# ********************************************************
def _get_intermediate_layers_not_chunked(self, x, n=1):
x = self.prepare_tokens_with_masks(x)
# If n is an int, take the n last blocks. If it's a list, take them
output, total_block_len = [], len(self.blocks)
blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
for i, blk in enumerate(self.blocks):
x = blk(x)
if i in blocks_to_take:
output.append(x)
assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
return output
def _get_intermediate_layers_chunked(self, x, n=1):
x = self.prepare_tokens_with_masks(x)
output, i, total_block_len = [], 0, len(self.blocks[-1])
# If n is an int, take the n last blocks. If it's a list, take them
blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
for block_chunk in self.blocks:
for blk in block_chunk[i:]: # Passing the nn.Identity()
x = blk(x)
if i in blocks_to_take:
output.append(x)
i += 1
assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
return output
def get_intermediate_layers(
self,
x: torch.Tensor,
n: Union[int, Sequence] = 1, # Layers or n last layers to take
reshape: bool = False,
return_class_token: bool = False,
norm=True,
) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
if self.chunked_blocks:
outputs = self._get_intermediate_layers_chunked(x, n)
else:
outputs = self._get_intermediate_layers_not_chunked(x, n)
if norm:
outputs = [self.norm(out) for out in outputs]
class_tokens = [out[:, 0] for out in outputs]
outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]
if reshape:
B, _, w, h = x.shape
outputs = [
out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
for out in outputs
]
if return_class_token:
return tuple(zip(outputs, class_tokens))
return tuple(outputs)
def forward(self, *args, is_training=False, **kwargs):
ret = self.forward_features(*args, **kwargs)
if is_training:
return ret
else:
return self.head(ret["x_norm_clstoken"])
def init_weights_vit_timm(module: nn.Module, name: str = ""):
"""ViT weight initialization, original timm impl (for reproducibility)"""
if isinstance(module, nn.Linear):
trunc_normal_(module.weight, std=0.02)
if module.bias is not None:
nn.init.zeros_(module.bias)
# ********** Modified by Zexin He in 2023-2024 **********
# block class selected from Block and BlockWithModulation
def _block_cls(**kwargs):
modulation_dim = kwargs.get("modulation_dim", None)
if modulation_dim is None:
block_cls = Block
else:
block_cls = BlockWithModulation
return block_cls
def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
model = DinoVisionTransformer(
patch_size=patch_size,
embed_dim=384,
depth=12,
num_heads=6,
mlp_ratio=4,
block_fn=partial(_block_cls(**kwargs), attn_class=MemEffAttention),
num_register_tokens=num_register_tokens,
**kwargs,
)
return model
def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
model = DinoVisionTransformer(
patch_size=patch_size,
embed_dim=768,
depth=12,
num_heads=12,
mlp_ratio=4,
block_fn=partial(_block_cls(**kwargs), attn_class=MemEffAttention),
num_register_tokens=num_register_tokens,
**kwargs,
)
return model
def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
model = DinoVisionTransformer(
patch_size=patch_size,
embed_dim=1024,
depth=24,
num_heads=16,
mlp_ratio=4,
block_fn=partial(_block_cls(**kwargs), attn_class=MemEffAttention),
num_register_tokens=num_register_tokens,
**kwargs,
)
return model
def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
"""
Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
"""
model = DinoVisionTransformer(
patch_size=patch_size,
embed_dim=1536,
depth=40,
num_heads=24,
mlp_ratio=4,
block_fn=partial(_block_cls(**kwargs), attn_class=MemEffAttention),
num_register_tokens=num_register_tokens,
**kwargs,
)
return model
# ********************************************************
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
import logging
from . import vision_transformer as vits
logger = logging.getLogger("dinov2")
def build_model(args, only_teacher=False, img_size=224):
args.arch = args.arch.removesuffix("_memeff")
if "vit" in args.arch:
vit_kwargs = dict(
img_size=img_size,
patch_size=args.patch_size,
init_values=args.layerscale,
ffn_layer=args.ffn_layer,
block_chunks=args.block_chunks,
qkv_bias=args.qkv_bias,
proj_bias=args.proj_bias,
ffn_bias=args.ffn_bias,
num_register_tokens=args.num_register_tokens,
interpolate_offset=args.interpolate_offset,
interpolate_antialias=args.interpolate_antialias,
)
teacher = vits.__dict__[args.arch](**vit_kwargs)
if only_teacher:
return teacher, teacher.embed_dim
student = vits.__dict__[args.arch](
**vit_kwargs,
drop_path_rate=args.drop_path_rate,
drop_path_uniform=args.drop_path_uniform,
)
embed_dim = student.embed_dim
return student, teacher, embed_dim
def build_model_from_cfg(cfg, only_teacher=False):
return build_model(cfg.student, only_teacher=only_teacher, img_size=cfg.crops.global_crops_size)
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
# References:
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
import logging
import os
import warnings
from torch import Tensor
from torch import nn
logger = logging.getLogger("dinov2")
XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
try:
if XFORMERS_ENABLED:
from xformers.ops import memory_efficient_attention, unbind
XFORMERS_AVAILABLE = True
warnings.warn("xFormers is available (Attention)")
else:
warnings.warn("xFormers is disabled (Attention)")
raise ImportError
except ImportError:
XFORMERS_AVAILABLE = False
warnings.warn("xFormers is not available (Attention)")
class Attention(nn.Module):
def __init__(
self,
dim: int,
num_heads: int = 8,
qkv_bias: bool = False,
proj_bias: bool = True,
attn_drop: float = 0.0,
proj_drop: float = 0.0,
) -> None:
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = head_dim**-0.5
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim, bias=proj_bias)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x: Tensor) -> Tensor:
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
attn = q @ k.transpose(-2, -1)
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
class MemEffAttention(Attention):
def forward(self, x: Tensor, attn_bias=None) -> Tensor:
if not XFORMERS_AVAILABLE:
if attn_bias is not None:
raise AssertionError("xFormers is required for using nested tensors")
return super().forward(x)
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
q, k, v = unbind(qkv, 2)
x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
x = x.reshape([B, N, C])
x = self.proj(x)
x = self.proj_drop(x)
return x
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
from typing import Union
import torch
from torch import Tensor
from torch import nn
class LayerScale(nn.Module):
def __init__(
self,
dim: int,
init_values: Union[float, Tensor] = 1e-5,
inplace: bool = False,
) -> None:
super().__init__()
self.inplace = inplace
self.gamma = nn.Parameter(init_values * torch.ones(dim))
def forward(self, x: Tensor) -> Tensor:
return x.mul_(self.gamma) if self.inplace else x * self.gamma
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
import torch
import torch.nn as nn
from torch.nn.init import trunc_normal_
from torch.nn.utils import weight_norm
class DINOHead(nn.Module):
def __init__(
self,
in_dim,
out_dim,
use_bn=False,
nlayers=3,
hidden_dim=2048,
bottleneck_dim=256,
mlp_bias=True,
):
super().__init__()
nlayers = max(nlayers, 1)
self.mlp = _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=hidden_dim, use_bn=use_bn, bias=mlp_bias)
self.apply(self._init_weights)
self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
self.last_layer.weight_g.data.fill_(1)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=0.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
def forward(self, x):
x = self.mlp(x)
eps = 1e-6 if x.dtype == torch.float16 else 1e-12
x = nn.functional.normalize(x, dim=-1, p=2, eps=eps)
x = self.last_layer(x)
return x
def _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True):
if nlayers == 1:
return nn.Linear(in_dim, bottleneck_dim, bias=bias)
else:
layers = [nn.Linear(in_dim, hidden_dim, bias=bias)]
if use_bn:
layers.append(nn.BatchNorm1d(hidden_dim))
layers.append(nn.GELU())
for _ in range(nlayers - 2):
layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias))
if use_bn:
layers.append(nn.BatchNorm1d(hidden_dim))
layers.append(nn.GELU())
layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias))
return nn.Sequential(*layers)
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
# References:
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
from typing import Callable, Optional
from torch import Tensor, nn
class Mlp(nn.Module):
def __init__(
self,
in_features: int,
hidden_features: Optional[int] = None,
out_features: Optional[int] = None,
act_layer: Callable[..., nn.Module] = nn.GELU,
drop: float = 0.0,
bias: bool = True,
) -> None:
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
self.drop = nn.Dropout(drop)
def forward(self, x: Tensor) -> Tensor:
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
# References:
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
from typing import Callable, Optional, Tuple, Union
from torch import Tensor
import torch.nn as nn
def make_2tuple(x):
if isinstance(x, tuple):
assert len(x) == 2
return x
assert isinstance(x, int)
return (x, x)
class PatchEmbed(nn.Module):
"""
2D image to patch embedding: (B,C,H,W) -> (B,N,D)
Args:
img_size: Image size.
patch_size: Patch token size.
in_chans: Number of input image channels.
embed_dim: Number of linear projection output channels.
norm_layer: Normalization layer.
"""
def __init__(
self,
img_size: Union[int, Tuple[int, int]] = 224,
patch_size: Union[int, Tuple[int, int]] = 16,
in_chans: int = 3,
embed_dim: int = 768,
norm_layer: Optional[Callable] = None,
flatten_embedding: bool = True,
) -> None:
super().__init__()
image_HW = make_2tuple(img_size)
patch_HW = make_2tuple(patch_size)
patch_grid_size = (
image_HW[0] // patch_HW[0],
image_HW[1] // patch_HW[1],
)
self.img_size = image_HW
self.patch_size = patch_HW
self.patches_resolution = patch_grid_size
self.num_patches = patch_grid_size[0] * patch_grid_size[1]
self.in_chans = in_chans
self.embed_dim = embed_dim
self.flatten_embedding = flatten_embedding
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
def forward(self, x: Tensor) -> Tensor:
_, _, H, W = x.shape
patch_H, patch_W = self.patch_size
assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
x = self.proj(x) # B C H W
H, W = x.size(2), x.size(3)
x = x.flatten(2).transpose(1, 2) # B HW C
x = self.norm(x)
if not self.flatten_embedding:
x = x.reshape(-1, H, W, self.embed_dim) # B H W C
return x
def flops(self) -> float:
Ho, Wo = self.patches_resolution
flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
if self.norm is not None:
flops += Ho * Wo * self.embed_dim
return flops
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
import os
from typing import Callable, Optional
import warnings
from torch import Tensor, nn
import torch.nn.functional as F
class SwiGLUFFN(nn.Module):
def __init__(
self,
in_features: int,
hidden_features: Optional[int] = None,
out_features: Optional[int] = None,
act_layer: Callable[..., nn.Module] = None,
drop: float = 0.0,
bias: bool = True,
) -> None:
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
def forward(self, x: Tensor) -> Tensor:
x12 = self.w12(x)
x1, x2 = x12.chunk(2, dim=-1)
hidden = F.silu(x1) * x2
return self.w3(hidden)
XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
try:
if XFORMERS_ENABLED:
from xformers.ops import SwiGLU
XFORMERS_AVAILABLE = True
warnings.warn("xFormers is available (SwiGLU)")
else:
warnings.warn("xFormers is disabled (SwiGLU)")
raise ImportError
except ImportError:
SwiGLU = SwiGLUFFN
XFORMERS_AVAILABLE = False
warnings.warn("xFormers is not available (SwiGLU)")
class SwiGLUFFNFused(SwiGLU):
def __init__(
self,
in_features: int,
hidden_features: Optional[int] = None,
out_features: Optional[int] = None,
act_layer: Callable[..., nn.Module] = None,
drop: float = 0.0,
bias: bool = True,
) -> None:
out_features = out_features or in_features
hidden_features = hidden_features or in_features
hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
super().__init__(
in_features=in_features,
hidden_features=hidden_features,
out_features=out_features,
bias=bias,
)
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
# References:
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
# ******************************************************************************
# Code modified by Zexin He in 2023-2024.
# Modifications are marked with clearly visible comments
# licensed under the Apache License, Version 2.0.
# ******************************************************************************
import logging
import os
from typing import Callable, List, Any, Tuple, Dict
import warnings
import torch
from torch import nn, Tensor
from .attention import Attention, MemEffAttention
from .drop_path import DropPath
from .layer_scale import LayerScale
from .mlp import Mlp
logger = logging.getLogger("dinov2")
XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
try:
if XFORMERS_ENABLED:
from xformers.ops import fmha, scaled_index_add, index_select_cat
XFORMERS_AVAILABLE = True
warnings.warn("xFormers is available (Block)")
else:
warnings.warn("xFormers is disabled (Block)")
raise ImportError
except ImportError:
XFORMERS_AVAILABLE = False
warnings.warn("xFormers is not available (Block)")
class Block(nn.Module):
def __init__(
self,
dim: int,
num_heads: int,
mlp_ratio: float = 4.0,
qkv_bias: bool = False,
proj_bias: bool = True,
ffn_bias: bool = True,
drop: float = 0.0,
attn_drop: float = 0.0,
init_values=None,
drop_path: float = 0.0,
act_layer: Callable[..., nn.Module] = nn.GELU,
norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
attn_class: Callable[..., nn.Module] = Attention,
ffn_layer: Callable[..., nn.Module] = Mlp,
) -> None:
super().__init__()
# print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
self.norm1 = norm_layer(dim)
self.attn = attn_class(
dim,
num_heads=num_heads,
qkv_bias=qkv_bias,
proj_bias=proj_bias,
attn_drop=attn_drop,
proj_drop=drop,
)
self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = ffn_layer(
in_features=dim,
hidden_features=mlp_hidden_dim,
act_layer=act_layer,
drop=drop,
bias=ffn_bias,
)
self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
self.sample_drop_ratio = drop_path
def forward(self, x: Tensor) -> Tensor:
def attn_residual_func(x: Tensor) -> Tensor:
return self.ls1(self.attn(self.norm1(x)))
def ffn_residual_func(x: Tensor) -> Tensor:
return self.ls2(self.mlp(self.norm2(x)))
if self.training and self.sample_drop_ratio > 0.1:
# the overhead is compensated only for a drop path rate larger than 0.1
x = drop_add_residual_stochastic_depth(
x,
residual_func=attn_residual_func,
sample_drop_ratio=self.sample_drop_ratio,
)
x = drop_add_residual_stochastic_depth(
x,
residual_func=ffn_residual_func,
sample_drop_ratio=self.sample_drop_ratio,
)
elif self.training and self.sample_drop_ratio > 0.0:
x = x + self.drop_path1(attn_residual_func(x))
x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2
else:
x = x + attn_residual_func(x)
x = x + ffn_residual_func(x)
return x
# ********** Modified by Zexin He in 2023-2024 **********
# Override forward with modulation input
class BlockWithModulation(Block):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def forward(self, x: Tensor, mod: Tensor) -> Tensor:
def attn_residual_func(x: Tensor, mod: Tensor) -> Tensor:
return self.ls1(self.attn(self.norm1(x, mod)))
def ffn_residual_func(x: Tensor, mod: Tensor) -> Tensor:
return self.ls2(self.mlp(self.norm2(x, mod)))
if self.training and self.sample_drop_ratio > 0.1:
raise NotImplementedError("Modulation with drop path ratio larger than 0.1 is not supported yet")
elif self.training and self.sample_drop_ratio > 0.0:
x = x + self.drop_path1(attn_residual_func(x, mod))
x = x + self.drop_path1(ffn_residual_func(x, mod)) # FIXME: drop_path2
else:
x = x + attn_residual_func(x, mod)
x = x + ffn_residual_func(x, mod)
return x
# ********************************************************
def drop_add_residual_stochastic_depth(
x: Tensor,
residual_func: Callable[[Tensor], Tensor],
sample_drop_ratio: float = 0.0,
) -> Tensor:
# 1) extract subset using permutation
b, n, d = x.shape
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
x_subset = x[brange]
# 2) apply residual_func to get residual
residual = residual_func(x_subset)
x_flat = x.flatten(1)
residual = residual.flatten(1)
residual_scale_factor = b / sample_subset_size
# 3) add the residual
x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
return x_plus_residual.view_as(x)
def get_branges_scales(x, sample_drop_ratio=0.0):
b, n, d = x.shape
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
residual_scale_factor = b / sample_subset_size
return brange, residual_scale_factor
def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
if scaling_vector is None:
x_flat = x.flatten(1)
residual = residual.flatten(1)
x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
else:
x_plus_residual = scaled_index_add(
x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
)
return x_plus_residual
attn_bias_cache: Dict[Tuple, Any] = {}
def get_attn_bias_and_cat(x_list, branges=None):
"""
this will perform the index select, cat the tensors, and provide the attn_bias from cache
"""
batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
if all_shapes not in attn_bias_cache.keys():
seqlens = []
for b, x in zip(batch_sizes, x_list):
for _ in range(b):
seqlens.append(x.shape[1])
attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
attn_bias._batch_sizes = batch_sizes
attn_bias_cache[all_shapes] = attn_bias
if branges is not None:
cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
else:
tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
cat_tensors = torch.cat(tensors_bs1, dim=1)
return attn_bias_cache[all_shapes], cat_tensors
def drop_add_residual_stochastic_depth_list(
x_list: List[Tensor],
residual_func: Callable[[Tensor, Any], Tensor],
sample_drop_ratio: float = 0.0,
scaling_vector=None,
) -> Tensor:
# 1) generate random set of indices for dropping samples in the batch
branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
branges = [s[0] for s in branges_scales]
residual_scale_factors = [s[1] for s in branges_scales]
# 2) get attention bias and index+concat the tensors
attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
# 3) apply residual_func to get residual, and split the result
residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore
outputs = []
for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
return outputs
class NestedTensorBlock(Block):
# ********** Modified by Zexin He in 2023-2024 **********
warnings.warn("NestedTensorBlock is deprecated for now!", DeprecationWarning)
# ********************************************************
def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
"""
x_list contains a list of tensors to nest together and run
"""
assert isinstance(self.attn, MemEffAttention)
if self.training and self.sample_drop_ratio > 0.0:
def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
return self.attn(self.norm1(x), attn_bias=attn_bias)
def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
return self.mlp(self.norm2(x))
x_list = drop_add_residual_stochastic_depth_list(
x_list,
residual_func=attn_residual_func,
sample_drop_ratio=self.sample_drop_ratio,
scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
)
x_list = drop_add_residual_stochastic_depth_list(
x_list,
residual_func=ffn_residual_func,
sample_drop_ratio=self.sample_drop_ratio,
scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
)
return x_list
else:
def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
return self.ls2(self.mlp(self.norm2(x)))
attn_bias, x = get_attn_bias_and_cat(x_list)
x = x + attn_residual_func(x, attn_bias=attn_bias)
x = x + ffn_residual_func(x)
return attn_bias.split(x)
def forward(self, x_or_x_list):
if isinstance(x_or_x_list, Tensor):
return super().forward(x_or_x_list)
elif isinstance(x_or_x_list, list):
if not XFORMERS_AVAILABLE:
raise AssertionError("xFormers is required for using nested tensors")
return self.forward_nested(x_or_x_list)
else:
raise AssertionError
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
# ******************************************************************************
# Code modified by Zexin He in 2023-2024.
# Modifications are marked with clearly visible comments
# licensed under the Apache License, Version 2.0.
# ******************************************************************************
from .dino_head import DINOHead
from .mlp import Mlp
from .patch_embed import PatchEmbed
from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
# ********** Modified by Zexin He in 2023-2024 **********
# Avoid using nested tensor for now, deprecating usage of NestedTensorBlock
from .block import Block, BlockWithModulation
# ********************************************************
from .attention import MemEffAttention
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
# References:
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
from torch import nn
def drop_path(x, drop_prob: float = 0.0, training: bool = False):
if drop_prob == 0.0 or not training:
return x
keep_prob = 1 - drop_prob
shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
if keep_prob > 0.0:
random_tensor.div_(keep_prob)
output = x * random_tensor
return output
class DropPath(nn.Module):
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
def __init__(self, drop_prob=None):
super(DropPath, self).__init__()
self.drop_prob = drop_prob
def forward(self, x):
return drop_path(x, self.drop_prob, self.training)
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
import itertools
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
_DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2"
def _make_dinov2_model_name(arch_name: str, patch_size: int, num_register_tokens: int = 0) -> str:
compact_arch_name = arch_name.replace("_", "")[:4]
registers_suffix = f"_reg{num_register_tokens}" if num_register_tokens else ""
return f"dinov2_{compact_arch_name}{patch_size}{registers_suffix}"
class CenterPadding(nn.Module):
def __init__(self, multiple):
super().__init__()
self.multiple = multiple
def _get_pad(self, size):
new_size = math.ceil(size / self.multiple) * self.multiple
pad_size = new_size - size
pad_size_left = pad_size // 2
pad_size_right = pad_size - pad_size_left
return pad_size_left, pad_size_right
@torch.inference_mode()
def forward(self, x):
pads = list(itertools.chain.from_iterable(self._get_pad(m) for m in x.shape[:1:-1]))
output = F.pad(x, pads)
return output
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
from enum import Enum
from functools import partial
from typing import Optional, Tuple, Union
import torch
from .backbones import _make_dinov2_model
from .depth import BNHead, DepthEncoderDecoder, DPTHead
from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name, CenterPadding
class Weights(Enum):
NYU = "NYU"
KITTI = "KITTI"
def _get_depth_range(pretrained: bool, weights: Weights = Weights.NYU) -> Tuple[float, float]:
if not pretrained: # Default
return (0.001, 10.0)
# Pretrained, set according to the training dataset for the provided weights
if weights == Weights.KITTI:
return (0.001, 80.0)
if weights == Weights.NYU:
return (0.001, 10.0)
return (0.001, 10.0)
def _make_dinov2_linear_depth_head(
*,
embed_dim: int,
layers: int,
min_depth: float,
max_depth: float,
**kwargs,
):
if layers not in (1, 4):
raise AssertionError(f"Unsupported number of layers: {layers}")
if layers == 1:
in_index = [0]
else:
assert layers == 4
in_index = [0, 1, 2, 3]
return BNHead(
classify=True,
n_bins=256,
bins_strategy="UD",
norm_strategy="linear",
upsample=4,
in_channels=[embed_dim] * len(in_index),
in_index=in_index,
input_transform="resize_concat",
channels=embed_dim * len(in_index) * 2,
align_corners=False,
min_depth=0.001,
max_depth=80,
loss_decode=(),
)
def _make_dinov2_linear_depther(
*,
arch_name: str = "vit_large",
layers: int = 4,
pretrained: bool = True,
weights: Union[Weights, str] = Weights.NYU,
depth_range: Optional[Tuple[float, float]] = None,
**kwargs,
):
if layers not in (1, 4):
raise AssertionError(f"Unsupported number of layers: {layers}")
if isinstance(weights, str):
try:
weights = Weights[weights]
except KeyError:
raise AssertionError(f"Unsupported weights: {weights}")
if depth_range is None:
depth_range = _get_depth_range(pretrained, weights)
min_depth, max_depth = depth_range
backbone = _make_dinov2_model(arch_name=arch_name, pretrained=pretrained, **kwargs)
embed_dim = backbone.embed_dim
patch_size = backbone.patch_size
model_name = _make_dinov2_model_name(arch_name, patch_size)
linear_depth_head = _make_dinov2_linear_depth_head(
embed_dim=embed_dim,
layers=layers,
min_depth=min_depth,
max_depth=max_depth,
)
layer_count = {
"vit_small": 12,
"vit_base": 12,
"vit_large": 24,
"vit_giant2": 40,
}[arch_name]
if layers == 4:
out_index = {
"vit_small": [2, 5, 8, 11],
"vit_base": [2, 5, 8, 11],
"vit_large": [4, 11, 17, 23],
"vit_giant2": [9, 19, 29, 39],
}[arch_name]
else:
assert layers == 1
out_index = [layer_count - 1]
model = DepthEncoderDecoder(backbone=backbone, decode_head=linear_depth_head)
model.backbone.forward = partial(
backbone.get_intermediate_layers,
n=out_index,
reshape=True,
return_class_token=True,
norm=False,
)
model.backbone.register_forward_pre_hook(lambda _, x: CenterPadding(patch_size)(x[0]))
if pretrained:
layers_str = str(layers) if layers == 4 else ""
weights_str = weights.value.lower()
url = _DINOV2_BASE_URL + f"/{model_name}/{model_name}_{weights_str}_linear{layers_str}_head.pth"
checkpoint = torch.hub.load_state_dict_from_url(url, map_location="cpu")
if "state_dict" in checkpoint:
state_dict = checkpoint["state_dict"]
model.load_state_dict(state_dict, strict=False)
return model
def dinov2_vits14_ld(*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
return _make_dinov2_linear_depther(
arch_name="vit_small", layers=layers, pretrained=pretrained, weights=weights, **kwargs
)
def dinov2_vitb14_ld(*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
return _make_dinov2_linear_depther(
arch_name="vit_base", layers=layers, pretrained=pretrained, weights=weights, **kwargs
)
def dinov2_vitl14_ld(*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
return _make_dinov2_linear_depther(
arch_name="vit_large", layers=layers, pretrained=pretrained, weights=weights, **kwargs
)
def dinov2_vitg14_ld(*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
return _make_dinov2_linear_depther(
arch_name="vit_giant2", layers=layers, ffn_layer="swiglufused", pretrained=pretrained, weights=weights, **kwargs
)
def _make_dinov2_dpt_depth_head(*, embed_dim: int, min_depth: float, max_depth: float):
return DPTHead(
in_channels=[embed_dim] * 4,
channels=256,
embed_dims=embed_dim,
post_process_channels=[embed_dim // 2 ** (3 - i) for i in range(4)],
readout_type="project",
min_depth=min_depth,
max_depth=max_depth,
loss_decode=(),
)
def _make_dinov2_dpt_depther(
*,
arch_name: str = "vit_large",
pretrained: bool = True,
weights: Union[Weights, str] = Weights.NYU,
depth_range: Optional[Tuple[float, float]] = None,
**kwargs,
):
if isinstance(weights, str):
try:
weights = Weights[weights]
except KeyError:
raise AssertionError(f"Unsupported weights: {weights}")
if depth_range is None:
depth_range = _get_depth_range(pretrained, weights)
min_depth, max_depth = depth_range
backbone = _make_dinov2_model(arch_name=arch_name, pretrained=pretrained, **kwargs)
model_name = _make_dinov2_model_name(arch_name, backbone.patch_size)
dpt_depth_head = _make_dinov2_dpt_depth_head(embed_dim=backbone.embed_dim, min_depth=min_depth, max_depth=max_depth)
out_index = {
"vit_small": [2, 5, 8, 11],
"vit_base": [2, 5, 8, 11],
"vit_large": [4, 11, 17, 23],
"vit_giant2": [9, 19, 29, 39],
}[arch_name]
model = DepthEncoderDecoder(backbone=backbone, decode_head=dpt_depth_head)
model.backbone.forward = partial(
backbone.get_intermediate_layers,
n=out_index,
reshape=True,
return_class_token=True,
norm=False,
)
model.backbone.register_forward_pre_hook(lambda _, x: CenterPadding(backbone.patch_size)(x[0]))
if pretrained:
weights_str = weights.value.lower()
url = _DINOV2_BASE_URL + f"/{model_name}/{model_name}_{weights_str}_dpt_head.pth"
checkpoint = torch.hub.load_state_dict_from_url(url, map_location="cpu")
if "state_dict" in checkpoint:
state_dict = checkpoint["state_dict"]
model.load_state_dict(state_dict, strict=False)
return model
def dinov2_vits14_dd(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
return _make_dinov2_dpt_depther(arch_name="vit_small", pretrained=pretrained, weights=weights, **kwargs)
def dinov2_vitb14_dd(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
return _make_dinov2_dpt_depther(arch_name="vit_base", pretrained=pretrained, weights=weights, **kwargs)
def dinov2_vitl14_dd(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
return _make_dinov2_dpt_depther(arch_name="vit_large", pretrained=pretrained, weights=weights, **kwargs)
def dinov2_vitg14_dd(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
return _make_dinov2_dpt_depther(
arch_name="vit_giant2", ffn_layer="swiglufused", pretrained=pretrained, weights=weights, **kwargs
)
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
from enum import Enum
from typing import Union
import torch
import torch.nn as nn
from .backbones import _make_dinov2_model
from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name
class Weights(Enum):
IMAGENET1K = "IMAGENET1K"
def _make_dinov2_linear_classification_head(
*,
arch_name: str = "vit_large",
patch_size: int = 14,
embed_dim: int = 1024,
layers: int = 4,
pretrained: bool = True,
weights: Union[Weights, str] = Weights.IMAGENET1K,
num_register_tokens: int = 0,
**kwargs,
):
if layers not in (1, 4):
raise AssertionError(f"Unsupported number of layers: {layers}")
if isinstance(weights, str):
try:
weights = Weights[weights]
except KeyError:
raise AssertionError(f"Unsupported weights: {weights}")
linear_head = nn.Linear((1 + layers) * embed_dim, 1_000)
if pretrained:
model_base_name = _make_dinov2_model_name(arch_name, patch_size)
model_full_name = _make_dinov2_model_name(arch_name, patch_size, num_register_tokens)
layers_str = str(layers) if layers == 4 else ""
url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_linear{layers_str}_head.pth"
state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu")
linear_head.load_state_dict(state_dict, strict=True)
return linear_head
class _LinearClassifierWrapper(nn.Module):
def __init__(self, *, backbone: nn.Module, linear_head: nn.Module, layers: int = 4):
super().__init__()
self.backbone = backbone
self.linear_head = linear_head
self.layers = layers
def forward(self, x):
if self.layers == 1:
x = self.backbone.forward_features(x)
cls_token = x["x_norm_clstoken"]
patch_tokens = x["x_norm_patchtokens"]
# fmt: off
linear_input = torch.cat([
cls_token,
patch_tokens.mean(dim=1),
], dim=1)
# fmt: on
elif self.layers == 4:
x = self.backbone.get_intermediate_layers(x, n=4, return_class_token=True)
# fmt: off
linear_input = torch.cat([
x[0][1],
x[1][1],
x[2][1],
x[3][1],
x[3][0].mean(dim=1),
], dim=1)
# fmt: on
else:
assert False, f"Unsupported number of layers: {self.layers}"
return self.linear_head(linear_input)
def _make_dinov2_linear_classifier(
*,
arch_name: str = "vit_large",
layers: int = 4,
pretrained: bool = True,
weights: Union[Weights, str] = Weights.IMAGENET1K,
num_register_tokens: int = 0,
interpolate_antialias: bool = False,
interpolate_offset: float = 0.1,
**kwargs,
):
backbone = _make_dinov2_model(
arch_name=arch_name,
pretrained=pretrained,
num_register_tokens=num_register_tokens,
interpolate_antialias=interpolate_antialias,
interpolate_offset=interpolate_offset,
**kwargs,
)
embed_dim = backbone.embed_dim
patch_size = backbone.patch_size
linear_head = _make_dinov2_linear_classification_head(
arch_name=arch_name,
patch_size=patch_size,
embed_dim=embed_dim,
layers=layers,
pretrained=pretrained,
weights=weights,
num_register_tokens=num_register_tokens,
)
return _LinearClassifierWrapper(backbone=backbone, linear_head=linear_head, layers=layers)
def dinov2_vits14_lc(
*,
layers: int = 4,
pretrained: bool = True,
weights: Union[Weights, str] = Weights.IMAGENET1K,
**kwargs,
):
"""
Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-S/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
"""
return _make_dinov2_linear_classifier(
arch_name="vit_small",
layers=layers,
pretrained=pretrained,
weights=weights,
**kwargs,
)
def dinov2_vitb14_lc(
*,
layers: int = 4,
pretrained: bool = True,
weights: Union[Weights, str] = Weights.IMAGENET1K,
**kwargs,
):
"""
Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-B/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
"""
return _make_dinov2_linear_classifier(
arch_name="vit_base",
layers=layers,
pretrained=pretrained,
weights=weights,
**kwargs,
)
def dinov2_vitl14_lc(
*,
layers: int = 4,
pretrained: bool = True,
weights: Union[Weights, str] = Weights.IMAGENET1K,
**kwargs,
):
"""
Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-L/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
"""
return _make_dinov2_linear_classifier(
arch_name="vit_large",
layers=layers,
pretrained=pretrained,
weights=weights,
**kwargs,
)
def dinov2_vitg14_lc(
*,
layers: int = 4,
pretrained: bool = True,
weights: Union[Weights, str] = Weights.IMAGENET1K,
**kwargs,
):
"""
Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-g/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
"""
return _make_dinov2_linear_classifier(
arch_name="vit_giant2",
layers=layers,
ffn_layer="swiglufused",
pretrained=pretrained,
weights=weights,
**kwargs,
)
def dinov2_vits14_reg_lc(
*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.IMAGENET1K, **kwargs
):
"""
Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-S/14 backbone with registers (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
"""
return _make_dinov2_linear_classifier(
arch_name="vit_small",
layers=layers,
pretrained=pretrained,
weights=weights,
num_register_tokens=4,
interpolate_antialias=True,
interpolate_offset=0.0,
**kwargs,
)
def dinov2_vitb14_reg_lc(
*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.IMAGENET1K, **kwargs
):
"""
Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-B/14 backbone with registers (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
"""
return _make_dinov2_linear_classifier(
arch_name="vit_base",
layers=layers,
pretrained=pretrained,
weights=weights,
num_register_tokens=4,
interpolate_antialias=True,
interpolate_offset=0.0,
**kwargs,
)
def dinov2_vitl14_reg_lc(
*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.IMAGENET1K, **kwargs
):
"""
Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-L/14 backbone with registers (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
"""
return _make_dinov2_linear_classifier(
arch_name="vit_large",
layers=layers,
pretrained=pretrained,
weights=weights,
num_register_tokens=4,
interpolate_antialias=True,
interpolate_offset=0.0,
**kwargs,
)
def dinov2_vitg14_reg_lc(
*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.IMAGENET1K, **kwargs
):
"""
Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-g/14 backbone with registers (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
"""
return _make_dinov2_linear_classifier(
arch_name="vit_giant2",
layers=layers,
ffn_layer="swiglufused",
pretrained=pretrained,
weights=weights,
num_register_tokens=4,
interpolate_antialias=True,
interpolate_offset=0.0,
**kwargs,
)
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
from enum import Enum
from typing import Union
import torch
from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name
class Weights(Enum):
LVD142M = "LVD142M"
def _make_dinov2_model(
*,
arch_name: str = "vit_large",
img_size: int = 518,
patch_size: int = 14,
init_values: float = 1.0,
ffn_layer: str = "mlp",
block_chunks: int = 0,
num_register_tokens: int = 0,
interpolate_antialias: bool = False,
interpolate_offset: float = 0.1,
pretrained: bool = True,
weights: Union[Weights, str] = Weights.LVD142M,
**kwargs,
):
from ..models import vision_transformer as vits
if isinstance(weights, str):
try:
weights = Weights[weights]
except KeyError:
raise AssertionError(f"Unsupported weights: {weights}")
model_base_name = _make_dinov2_model_name(arch_name, patch_size)
vit_kwargs = dict(
img_size=img_size,
patch_size=patch_size,
init_values=init_values,
ffn_layer=ffn_layer,
block_chunks=block_chunks,
num_register_tokens=num_register_tokens,
interpolate_antialias=interpolate_antialias,
interpolate_offset=interpolate_offset,
)
vit_kwargs.update(**kwargs)
model = vits.__dict__[arch_name](**vit_kwargs)
if pretrained:
model_full_name = _make_dinov2_model_name(arch_name, patch_size, num_register_tokens)
url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_pretrain.pth"
state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu")
# ********** Modified by Zexin He in 2023-2024 **********
state_dict = {k: v for k, v in state_dict.items() if 'mask_token' not in k} # DDP concern
if vit_kwargs.get("modulation_dim") is not None:
state_dict = {
k.replace('norm1', 'norm1.norm').replace('norm2', 'norm2.norm'): v
for k, v in state_dict.items()
}
model.load_state_dict(state_dict, strict=False)
else:
model.load_state_dict(state_dict, strict=True)
# ********************************************************
return model
def dinov2_vits14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
"""
DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset.
"""
return _make_dinov2_model(arch_name="vit_small", pretrained=pretrained, weights=weights, **kwargs)
def dinov2_vitb14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
"""
DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset.
"""
return _make_dinov2_model(arch_name="vit_base", pretrained=pretrained, weights=weights, **kwargs)
def dinov2_vitl14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
"""
DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset.
"""
return _make_dinov2_model(arch_name="vit_large", pretrained=pretrained, weights=weights, **kwargs)
def dinov2_vitg14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
"""
DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset.
"""
return _make_dinov2_model(
arch_name="vit_giant2",
ffn_layer="swiglufused",
weights=weights,
pretrained=pretrained,
**kwargs,
)
def dinov2_vits14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
"""
DINOv2 ViT-S/14 model with registers (optionally) pretrained on the LVD-142M dataset.
"""
return _make_dinov2_model(
arch_name="vit_small",
pretrained=pretrained,
weights=weights,
num_register_tokens=4,
interpolate_antialias=True,
interpolate_offset=0.0,
**kwargs,
)
def dinov2_vitb14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
"""
DINOv2 ViT-B/14 model with registers (optionally) pretrained on the LVD-142M dataset.
"""
return _make_dinov2_model(
arch_name="vit_base",
pretrained=pretrained,
weights=weights,
num_register_tokens=4,
interpolate_antialias=True,
interpolate_offset=0.0,
**kwargs,
)
def dinov2_vitl14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
"""
DINOv2 ViT-L/14 model with registers (optionally) pretrained on the LVD-142M dataset.
"""
return _make_dinov2_model(
arch_name="vit_large",
pretrained=pretrained,
weights=weights,
num_register_tokens=4,
interpolate_antialias=True,
interpolate_offset=0.0,
**kwargs,
)
def dinov2_vitg14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
"""
DINOv2 ViT-g/14 model with registers (optionally) pretrained on the LVD-142M dataset.
"""
return _make_dinov2_model(
arch_name="vit_giant2",
ffn_layer="swiglufused",
weights=weights,
pretrained=pretrained,
num_register_tokens=4,
interpolate_antialias=True,
interpolate_offset=0.0,
**kwargs,
)
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F
from .ops import resize
def add_prefix(inputs, prefix):
"""Add prefix for dict.
Args:
inputs (dict): The input dict with str keys.
prefix (str): The prefix to add.
Returns:
dict: The dict with keys updated with ``prefix``.
"""
outputs = dict()
for name, value in inputs.items():
outputs[f"{prefix}.{name}"] = value
return outputs
class DepthEncoderDecoder(nn.Module):
"""Encoder Decoder depther.
EncoderDecoder typically consists of backbone and decode_head.
"""
def __init__(self, backbone, decode_head):
super(DepthEncoderDecoder, self).__init__()
self.backbone = backbone
self.decode_head = decode_head
self.align_corners = self.decode_head.align_corners
def extract_feat(self, img):
"""Extract features from images."""
return self.backbone(img)
def encode_decode(self, img, img_metas, rescale=True, size=None):
"""Encode images with backbone and decode into a depth estimation
map of the same size as input."""
x = self.extract_feat(img)
out = self._decode_head_forward_test(x, img_metas)
# crop the pred depth to the certain range.
out = torch.clamp(out, min=self.decode_head.min_depth, max=self.decode_head.max_depth)
if rescale:
if size is None:
if img_metas is not None:
size = img_metas[0]["ori_shape"][:2]
else:
size = img.shape[2:]
out = resize(input=out, size=size, mode="bilinear", align_corners=self.align_corners)
return out
def _decode_head_forward_train(self, img, x, img_metas, depth_gt, **kwargs):
"""Run forward function and calculate loss for decode head in
training."""
losses = dict()
loss_decode = self.decode_head.forward_train(img, x, img_metas, depth_gt, **kwargs)
losses.update(add_prefix(loss_decode, "decode"))
return losses
def _decode_head_forward_test(self, x, img_metas):
"""Run forward function and calculate loss for decode head in
inference."""
depth_pred = self.decode_head.forward_test(x, img_metas)
return depth_pred
def forward_dummy(self, img):
"""Dummy forward function."""
depth = self.encode_decode(img, None)
return depth
def forward_train(self, img, img_metas, depth_gt, **kwargs):
"""Forward function for training.
Args:
img (Tensor): Input images.
img_metas (list[dict]): List of image info dict where each dict
has: 'img_shape', 'scale_factor', 'flip', and may also contain
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
For details on the values of these keys see
`depth/datasets/pipelines/formatting.py:Collect`.
depth_gt (Tensor): Depth gt
used if the architecture supports depth estimation task.
Returns:
dict[str, Tensor]: a dictionary of loss components
"""
x = self.extract_feat(img)
losses = dict()
# the last of x saves the info from neck
loss_decode = self._decode_head_forward_train(img, x, img_metas, depth_gt, **kwargs)
losses.update(loss_decode)
return losses
def whole_inference(self, img, img_meta, rescale, size=None):
"""Inference with full image."""
return self.encode_decode(img, img_meta, rescale, size=size)
def slide_inference(self, img, img_meta, rescale, stride, crop_size):
"""Inference by sliding-window with overlap.
If h_crop > h_img or w_crop > w_img, the small patch will be used to
decode without padding.
"""
h_stride, w_stride = stride
h_crop, w_crop = crop_size
batch_size, _, h_img, w_img = img.size()
h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
preds = img.new_zeros((batch_size, 1, h_img, w_img))
count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
for h_idx in range(h_grids):
for w_idx in range(w_grids):
y1 = h_idx * h_stride
x1 = w_idx * w_stride
y2 = min(y1 + h_crop, h_img)
x2 = min(x1 + w_crop, w_img)
y1 = max(y2 - h_crop, 0)
x1 = max(x2 - w_crop, 0)
crop_img = img[:, :, y1:y2, x1:x2]
depth_pred = self.encode_decode(crop_img, img_meta, rescale)
preds += F.pad(depth_pred, (int(x1), int(preds.shape[3] - x2), int(y1), int(preds.shape[2] - y2)))
count_mat[:, :, y1:y2, x1:x2] += 1
assert (count_mat == 0).sum() == 0
if torch.onnx.is_in_onnx_export():
# cast count_mat to constant while exporting to ONNX
count_mat = torch.from_numpy(count_mat.cpu().detach().numpy()).to(device=img.device)
preds = preds / count_mat
return preds
def inference(self, img, img_meta, rescale, size=None, mode="whole"):
"""Inference with slide/whole style.
Args:
img (Tensor): The input image of shape (N, 3, H, W).
img_meta (dict): Image info dict where each dict has: 'img_shape',
'scale_factor', 'flip', and may also contain
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
For details on the values of these keys see
`depth/datasets/pipelines/formatting.py:Collect`.
rescale (bool): Whether rescale back to original shape.
Returns:
Tensor: The output depth map.
"""
assert mode in ["slide", "whole"]
ori_shape = img_meta[0]["ori_shape"]
assert all(_["ori_shape"] == ori_shape for _ in img_meta)
if mode == "slide":
depth_pred = self.slide_inference(img, img_meta, rescale)
else:
depth_pred = self.whole_inference(img, img_meta, rescale, size=size)
output = depth_pred
flip = img_meta[0]["flip"]
if flip:
flip_direction = img_meta[0]["flip_direction"]
assert flip_direction in ["horizontal", "vertical"]
if flip_direction == "horizontal":
output = output.flip(dims=(3,))
elif flip_direction == "vertical":
output = output.flip(dims=(2,))
return output
def simple_test(self, img, img_meta, rescale=True):
"""Simple test with single image."""
depth_pred = self.inference(img, img_meta, rescale)
if torch.onnx.is_in_onnx_export():
# our inference backend only support 4D output
depth_pred = depth_pred.unsqueeze(0)
return depth_pred
depth_pred = depth_pred.cpu().numpy()
# unravel batch dim
depth_pred = list(depth_pred)
return depth_pred
def aug_test(self, imgs, img_metas, rescale=True):
"""Test with augmentations.
Only rescale=True is supported.
"""
# aug_test rescale all imgs back to ori_shape for now
assert rescale
# to save memory, we get augmented depth logit inplace
depth_pred = self.inference(imgs[0], img_metas[0], rescale)
for i in range(1, len(imgs)):
cur_depth_pred = self.inference(imgs[i], img_metas[i], rescale, size=depth_pred.shape[-2:])
depth_pred += cur_depth_pred
depth_pred /= len(imgs)
depth_pred = depth_pred.cpu().numpy()
# unravel batch dim
depth_pred = list(depth_pred)
return depth_pred
def forward_test(self, imgs, img_metas, **kwargs):
"""
Args:
imgs (List[Tensor]): the outer list indicates test-time
augmentations and inner Tensor should have a shape NxCxHxW,
which contains all images in the batch.
img_metas (List[List[dict]]): the outer list indicates test-time
augs (multiscale, flip, etc.) and the inner list indicates
images in a batch.
"""
for var, name in [(imgs, "imgs"), (img_metas, "img_metas")]:
if not isinstance(var, list):
raise TypeError(f"{name} must be a list, but got " f"{type(var)}")
num_augs = len(imgs)
if num_augs != len(img_metas):
raise ValueError(f"num of augmentations ({len(imgs)}) != " f"num of image meta ({len(img_metas)})")
# all images in the same aug batch all of the same ori_shape and pad
# shape
for img_meta in img_metas:
ori_shapes = [_["ori_shape"] for _ in img_meta]
assert all(shape == ori_shapes[0] for shape in ori_shapes)
img_shapes = [_["img_shape"] for _ in img_meta]
assert all(shape == img_shapes[0] for shape in img_shapes)
pad_shapes = [_["pad_shape"] for _ in img_meta]
assert all(shape == pad_shapes[0] for shape in pad_shapes)
if num_augs == 1:
return self.simple_test(imgs[0], img_metas[0], **kwargs)
else:
return self.aug_test(imgs, img_metas, **kwargs)
def forward(self, img, img_metas, return_loss=True, **kwargs):
"""Calls either :func:`forward_train` or :func:`forward_test` depending
on whether ``return_loss`` is ``True``.
Note this setting will change the expected inputs. When
``return_loss=True``, img and img_meta are single-nested (i.e. Tensor
and List[dict]), and when ``resturn_loss=False``, img and img_meta
should be double nested (i.e. List[Tensor], List[List[dict]]), with
the outer list indicating test time augmentations.
"""
if return_loss:
return self.forward_train(img, img_metas, **kwargs)
else:
return self.forward_test(img, img_metas, **kwargs)
def train_step(self, data_batch, optimizer, **kwargs):
"""The iteration step during training.
This method defines an iteration step during training, except for the
back propagation and optimizer updating, which are done in an optimizer
hook. Note that in some complicated cases or models, the whole process
including back propagation and optimizer updating is also defined in
this method, such as GAN.
Args:
data (dict): The output of dataloader.
optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of
runner is passed to ``train_step()``. This argument is unused
and reserved.
Returns:
dict: It should contain at least 3 keys: ``loss``, ``log_vars``,
``num_samples``.
``loss`` is a tensor for back propagation, which can be a
weighted sum of multiple losses.
``log_vars`` contains all the variables to be sent to the
logger.
``num_samples`` indicates the batch size (when the model is
DDP, it means the batch size on each GPU), which is used for
averaging the logs.
"""
losses = self(**data_batch)
# split losses and images
real_losses = {}
log_imgs = {}
for k, v in losses.items():
if "img" in k:
log_imgs[k] = v
else:
real_losses[k] = v
loss, log_vars = self._parse_losses(real_losses)
outputs = dict(loss=loss, log_vars=log_vars, num_samples=len(data_batch["img_metas"]), log_imgs=log_imgs)
return outputs
def val_step(self, data_batch, **kwargs):
"""The iteration step during validation.
This method shares the same signature as :func:`train_step`, but used
during val epochs. Note that the evaluation after training epochs is
not implemented with this method, but an evaluation hook.
"""
output = self(**data_batch, **kwargs)
return output
@staticmethod
def _parse_losses(losses):
import torch.distributed as dist
"""Parse the raw outputs (losses) of the network.
Args:
losses (dict): Raw output of the network, which usually contain
losses and other necessary information.
Returns:
tuple[Tensor, dict]: (loss, log_vars), loss is the loss tensor
which may be a weighted sum of all losses, log_vars contains
all the variables to be sent to the logger.
"""
log_vars = OrderedDict()
for loss_name, loss_value in losses.items():
if isinstance(loss_value, torch.Tensor):
log_vars[loss_name] = loss_value.mean()
elif isinstance(loss_value, list):
log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)
else:
raise TypeError(f"{loss_name} is not a tensor or list of tensors")
loss = sum(_value for _key, _value in log_vars.items() if "loss" in _key)
log_vars["loss"] = loss
for loss_name, loss_value in log_vars.items():
# reduce loss when distributed training
if dist.is_available() and dist.is_initialized():
loss_value = loss_value.data.clone()
dist.all_reduce(loss_value.div_(dist.get_world_size()))
log_vars[loss_name] = loss_value.item()
return loss, log_vars
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
import copy
from functools import partial
import math
import warnings
import torch
import torch.nn as nn
from .ops import resize
# XXX: (Untested) replacement for mmcv.imdenormalize()
def _imdenormalize(img, mean, std, to_bgr=True):
import numpy as np
mean = mean.reshape(1, -1).astype(np.float64)
std = std.reshape(1, -1).astype(np.float64)
img = (img * std) + mean
if to_bgr:
img = img[::-1]
return img
class DepthBaseDecodeHead(nn.Module):
"""Base class for BaseDecodeHead.
Args:
in_channels (List): Input channels.
channels (int): Channels after modules, before conv_depth.
conv_layer (nn.Module): Conv layers. Default: None.
act_layer (nn.Module): Activation layers. Default: nn.ReLU.
loss_decode (dict): Config of decode loss.
Default: ().
sampler (dict|None): The config of depth map sampler.
Default: None.
align_corners (bool): align_corners argument of F.interpolate.
Default: False.
min_depth (int): Min depth in dataset setting.
Default: 1e-3.
max_depth (int): Max depth in dataset setting.
Default: None.
norm_layer (dict|None): Norm layers.
Default: None.
classify (bool): Whether predict depth in a cls.-reg. manner.
Default: False.
n_bins (int): The number of bins used in cls. step.
Default: 256.
bins_strategy (str): The discrete strategy used in cls. step.
Default: 'UD'.
norm_strategy (str): The norm strategy on cls. probability
distribution. Default: 'linear'
scale_up (str): Whether predict depth in a scale-up manner.
Default: False.
"""
def __init__(
self,
in_channels,
conv_layer=None,
act_layer=nn.ReLU,
channels=96,
loss_decode=(),
sampler=None,
align_corners=False,
min_depth=1e-3,
max_depth=None,
norm_layer=None,
classify=False,
n_bins=256,
bins_strategy="UD",
norm_strategy="linear",
scale_up=False,
):
super(DepthBaseDecodeHead, self).__init__()
self.in_channels = in_channels
self.channels = channels
self.conf_layer = conv_layer
self.act_layer = act_layer
self.loss_decode = loss_decode
self.align_corners = align_corners
self.min_depth = min_depth
self.max_depth = max_depth
self.norm_layer = norm_layer
self.classify = classify
self.n_bins = n_bins
self.scale_up = scale_up
if self.classify:
assert bins_strategy in ["UD", "SID"], "Support bins_strategy: UD, SID"
assert norm_strategy in ["linear", "softmax", "sigmoid"], "Support norm_strategy: linear, softmax, sigmoid"
self.bins_strategy = bins_strategy
self.norm_strategy = norm_strategy
self.softmax = nn.Softmax(dim=1)
self.conv_depth = nn.Conv2d(channels, n_bins, kernel_size=3, padding=1, stride=1)
else:
self.conv_depth = nn.Conv2d(channels, 1, kernel_size=3, padding=1, stride=1)
self.relu = nn.ReLU()
self.sigmoid = nn.Sigmoid()
def forward(self, inputs, img_metas):
"""Placeholder of forward function."""
pass
def forward_train(self, img, inputs, img_metas, depth_gt):
"""Forward function for training.
Args:
inputs (list[Tensor]): List of multi-level img features.
img_metas (list[dict]): List of image info dict where each dict
has: 'img_shape', 'scale_factor', 'flip', and may also contain
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
For details on the values of these keys see
`depth/datasets/pipelines/formatting.py:Collect`.
depth_gt (Tensor): GT depth
Returns:
dict[str, Tensor]: a dictionary of loss components
"""
depth_pred = self.forward(inputs, img_metas)
losses = self.losses(depth_pred, depth_gt)
log_imgs = self.log_images(img[0], depth_pred[0], depth_gt[0], img_metas[0])
losses.update(**log_imgs)
return losses
def forward_test(self, inputs, img_metas):
"""Forward function for testing.
Args:
inputs (list[Tensor]): List of multi-level img features.
img_metas (list[dict]): List of image info dict where each dict
has: 'img_shape', 'scale_factor', 'flip', and may also contain
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
For details on the values of these keys see
`depth/datasets/pipelines/formatting.py:Collect`.
Returns:
Tensor: Output depth map.
"""
return self.forward(inputs, img_metas)
def depth_pred(self, feat):
"""Prediction each pixel."""
if self.classify:
logit = self.conv_depth(feat)
if self.bins_strategy == "UD":
bins = torch.linspace(self.min_depth, self.max_depth, self.n_bins, device=feat.device)
elif self.bins_strategy == "SID":
bins = torch.logspace(self.min_depth, self.max_depth, self.n_bins, device=feat.device)
# following Adabins, default linear
if self.norm_strategy == "linear":
logit = torch.relu(logit)
eps = 0.1
logit = logit + eps
logit = logit / logit.sum(dim=1, keepdim=True)
elif self.norm_strategy == "softmax":
logit = torch.softmax(logit, dim=1)
elif self.norm_strategy == "sigmoid":
logit = torch.sigmoid(logit)
logit = logit / logit.sum(dim=1, keepdim=True)
output = torch.einsum("ikmn,k->imn", [logit, bins]).unsqueeze(dim=1)
else:
if self.scale_up:
output = self.sigmoid(self.conv_depth(feat)) * self.max_depth
else:
output = self.relu(self.conv_depth(feat)) + self.min_depth
return output
def losses(self, depth_pred, depth_gt):
"""Compute depth loss."""
loss = dict()
depth_pred = resize(
input=depth_pred, size=depth_gt.shape[2:], mode="bilinear", align_corners=self.align_corners, warning=False
)
if not isinstance(self.loss_decode, nn.ModuleList):
losses_decode = [self.loss_decode]
else:
losses_decode = self.loss_decode
for loss_decode in losses_decode:
if loss_decode.loss_name not in loss:
loss[loss_decode.loss_name] = loss_decode(depth_pred, depth_gt)
else:
loss[loss_decode.loss_name] += loss_decode(depth_pred, depth_gt)
return loss
def log_images(self, img_path, depth_pred, depth_gt, img_meta):
import numpy as np
show_img = copy.deepcopy(img_path.detach().cpu().permute(1, 2, 0))
show_img = show_img.numpy().astype(np.float32)
show_img = _imdenormalize(
show_img,
img_meta["img_norm_cfg"]["mean"],
img_meta["img_norm_cfg"]["std"],
img_meta["img_norm_cfg"]["to_rgb"],
)
show_img = np.clip(show_img, 0, 255)
show_img = show_img.astype(np.uint8)
show_img = show_img[:, :, ::-1]
show_img = show_img.transpose(0, 2, 1)
show_img = show_img.transpose(1, 0, 2)
depth_pred = depth_pred / torch.max(depth_pred)
depth_gt = depth_gt / torch.max(depth_gt)
depth_pred_color = copy.deepcopy(depth_pred.detach().cpu())
depth_gt_color = copy.deepcopy(depth_gt.detach().cpu())
return {"img_rgb": show_img, "img_depth_pred": depth_pred_color, "img_depth_gt": depth_gt_color}
class BNHead(DepthBaseDecodeHead):
"""Just a batchnorm."""
def __init__(self, input_transform="resize_concat", in_index=(0, 1, 2, 3), upsample=1, **kwargs):
super().__init__(**kwargs)
self.input_transform = input_transform
self.in_index = in_index
self.upsample = upsample
# self.bn = nn.SyncBatchNorm(self.in_channels)
if self.classify:
self.conv_depth = nn.Conv2d(self.channels, self.n_bins, kernel_size=1, padding=0, stride=1)
else:
self.conv_depth = nn.Conv2d(self.channels, 1, kernel_size=1, padding=0, stride=1)
def _transform_inputs(self, inputs):
"""Transform inputs for decoder.
Args:
inputs (list[Tensor]): List of multi-level img features.
Returns:
Tensor: The transformed inputs
"""
if "concat" in self.input_transform:
inputs = [inputs[i] for i in self.in_index]
if "resize" in self.input_transform:
inputs = [
resize(
input=x,
size=[s * self.upsample for s in inputs[0].shape[2:]],
mode="bilinear",
align_corners=self.align_corners,
)
for x in inputs
]
inputs = torch.cat(inputs, dim=1)
elif self.input_transform == "multiple_select":
inputs = [inputs[i] for i in self.in_index]
else:
inputs = inputs[self.in_index]
return inputs
def _forward_feature(self, inputs, img_metas=None, **kwargs):
"""Forward function for feature maps before classifying each pixel with
``self.cls_seg`` fc.
Args:
inputs (list[Tensor]): List of multi-level img features.
Returns:
feats (Tensor): A tensor of shape (batch_size, self.channels,
H, W) which is feature map for last layer of decoder head.
"""
# accept lists (for cls token)
inputs = list(inputs)
for i, x in enumerate(inputs):
if len(x) == 2:
x, cls_token = x[0], x[1]
if len(x.shape) == 2:
x = x[:, :, None, None]
cls_token = cls_token[:, :, None, None].expand_as(x)
inputs[i] = torch.cat((x, cls_token), 1)
else:
x = x[0]
if len(x.shape) == 2:
x = x[:, :, None, None]
inputs[i] = x
x = self._transform_inputs(inputs)
# feats = self.bn(x)
return x
def forward(self, inputs, img_metas=None, **kwargs):
"""Forward function."""
output = self._forward_feature(inputs, img_metas=img_metas, **kwargs)
output = self.depth_pred(output)
return output
class ConvModule(nn.Module):
"""A conv block that bundles conv/norm/activation layers.
This block simplifies the usage of convolution layers, which are commonly
used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).
It is based upon three build methods: `build_conv_layer()`,
`build_norm_layer()` and `build_activation_layer()`.
Besides, we add some additional features in this module.
1. Automatically set `bias` of the conv layer.
2. Spectral norm is supported.
3. More padding modes are supported. Before PyTorch 1.5, nn.Conv2d only
supports zero and circular padding, and we add "reflect" padding mode.
Args:
in_channels (int): Number of channels in the input feature map.
Same as that in ``nn._ConvNd``.
out_channels (int): Number of channels produced by the convolution.
Same as that in ``nn._ConvNd``.
kernel_size (int | tuple[int]): Size of the convolving kernel.
Same as that in ``nn._ConvNd``.
stride (int | tuple[int]): Stride of the convolution.
Same as that in ``nn._ConvNd``.
padding (int | tuple[int]): Zero-padding added to both sides of
the input. Same as that in ``nn._ConvNd``.
dilation (int | tuple[int]): Spacing between kernel elements.
Same as that in ``nn._ConvNd``.
groups (int): Number of blocked connections from input channels to
output channels. Same as that in ``nn._ConvNd``.
bias (bool | str): If specified as `auto`, it will be decided by the
norm_layer. Bias will be set as True if `norm_layer` is None, otherwise
False. Default: "auto".
conv_layer (nn.Module): Convolution layer. Default: None,
which means using conv2d.
norm_layer (nn.Module): Normalization layer. Default: None.
act_layer (nn.Module): Activation layer. Default: nn.ReLU.
inplace (bool): Whether to use inplace mode for activation.
Default: True.
with_spectral_norm (bool): Whether use spectral norm in conv module.
Default: False.
padding_mode (str): If the `padding_mode` has not been supported by
current `Conv2d` in PyTorch, we will use our own padding layer
instead. Currently, we support ['zeros', 'circular'] with official
implementation and ['reflect'] with our own implementation.
Default: 'zeros'.
order (tuple[str]): The order of conv/norm/activation layers. It is a
sequence of "conv", "norm" and "act". Common examples are
("conv", "norm", "act") and ("act", "conv", "norm").
Default: ('conv', 'norm', 'act').
"""
_abbr_ = "conv_block"
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
groups=1,
bias="auto",
conv_layer=nn.Conv2d,
norm_layer=None,
act_layer=nn.ReLU,
inplace=True,
with_spectral_norm=False,
padding_mode="zeros",
order=("conv", "norm", "act"),
):
super(ConvModule, self).__init__()
official_padding_mode = ["zeros", "circular"]
self.conv_layer = conv_layer
self.norm_layer = norm_layer
self.act_layer = act_layer
self.inplace = inplace
self.with_spectral_norm = with_spectral_norm
self.with_explicit_padding = padding_mode not in official_padding_mode
self.order = order
assert isinstance(self.order, tuple) and len(self.order) == 3
assert set(order) == set(["conv", "norm", "act"])
self.with_norm = norm_layer is not None
self.with_activation = act_layer is not None
# if the conv layer is before a norm layer, bias is unnecessary.
if bias == "auto":
bias = not self.with_norm
self.with_bias = bias
if self.with_explicit_padding:
if padding_mode == "zeros":
padding_layer = nn.ZeroPad2d
else:
raise AssertionError(f"Unsupported padding mode: {padding_mode}")
self.pad = padding_layer(padding)
# reset padding to 0 for conv module
conv_padding = 0 if self.with_explicit_padding else padding
# build convolution layer
self.conv = self.conv_layer(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=conv_padding,
dilation=dilation,
groups=groups,
bias=bias,
)
# export the attributes of self.conv to a higher level for convenience
self.in_channels = self.conv.in_channels
self.out_channels = self.conv.out_channels
self.kernel_size = self.conv.kernel_size
self.stride = self.conv.stride
self.padding = padding
self.dilation = self.conv.dilation
self.transposed = self.conv.transposed
self.output_padding = self.conv.output_padding
self.groups = self.conv.groups
if self.with_spectral_norm:
self.conv = nn.utils.spectral_norm(self.conv)
# build normalization layers
if self.with_norm:
# norm layer is after conv layer
if order.index("norm") > order.index("conv"):
norm_channels = out_channels
else:
norm_channels = in_channels
norm = partial(norm_layer, num_features=norm_channels)
self.add_module("norm", norm)
if self.with_bias:
from torch.nnModules.batchnorm import _BatchNorm
from torch.nnModules.instancenorm import _InstanceNorm
if isinstance(norm, (_BatchNorm, _InstanceNorm)):
warnings.warn("Unnecessary conv bias before batch/instance norm")
else:
self.norm_name = None
# build activation layer
if self.with_activation:
# nn.Tanh has no 'inplace' argument
# (nn.Tanh, nn.PReLU, nn.Sigmoid, nn.HSigmoid, nn.Swish, nn.GELU)
if not isinstance(act_layer, (nn.Tanh, nn.PReLU, nn.Sigmoid, nn.GELU)):
act_layer = partial(act_layer, inplace=inplace)
self.activate = act_layer()
# Use msra init by default
self.init_weights()
@property
def norm(self):
if self.norm_name:
return getattr(self, self.norm_name)
else:
return None
def init_weights(self):
# 1. It is mainly for customized conv layers with their own
# initialization manners by calling their own ``init_weights()``,
# and we do not want ConvModule to override the initialization.
# 2. For customized conv layers without their own initialization
# manners (that is, they don't have their own ``init_weights()``)
# and PyTorch's conv layers, they will be initialized by
# this method with default ``kaiming_init``.
# Note: For PyTorch's conv layers, they will be overwritten by our
# initialization implementation using default ``kaiming_init``.
if not hasattr(self.conv, "init_weights"):
if self.with_activation and isinstance(self.act_layer, nn.LeakyReLU):
nonlinearity = "leaky_relu"
a = 0.01 # XXX: default negative_slope
else:
nonlinearity = "relu"
a = 0
if hasattr(self.conv, "weight") and self.conv.weight is not None:
nn.init.kaiming_normal_(self.conv.weight, a=a, mode="fan_out", nonlinearity=nonlinearity)
if hasattr(self.conv, "bias") and self.conv.bias is not None:
nn.init.constant_(self.conv.bias, 0)
if self.with_norm:
if hasattr(self.norm, "weight") and self.norm.weight is not None:
nn.init.constant_(self.norm.weight, 1)
if hasattr(self.norm, "bias") and self.norm.bias is not None:
nn.init.constant_(self.norm.bias, 0)
def forward(self, x, activate=True, norm=True):
for layer in self.order:
if layer == "conv":
if self.with_explicit_padding:
x = self.pad(x)
x = self.conv(x)
elif layer == "norm" and norm and self.with_norm:
x = self.norm(x)
elif layer == "act" and activate and self.with_activation:
x = self.activate(x)
return x
class Interpolate(nn.Module):
def __init__(self, scale_factor, mode, align_corners=False):
super(Interpolate, self).__init__()
self.interp = nn.functional.interpolate
self.scale_factor = scale_factor
self.mode = mode
self.align_corners = align_corners
def forward(self, x):
x = self.interp(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners)
return x
class HeadDepth(nn.Module):
def __init__(self, features):
super(HeadDepth, self).__init__()
self.head = nn.Sequential(
nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
)
def forward(self, x):
x = self.head(x)
return x
class ReassembleBlocks(nn.Module):
"""ViTPostProcessBlock, process cls_token in ViT backbone output and
rearrange the feature vector to feature map.
Args:
in_channels (int): ViT feature channels. Default: 768.
out_channels (List): output channels of each stage.
Default: [96, 192, 384, 768].
readout_type (str): Type of readout operation. Default: 'ignore'.
patch_size (int): The patch size. Default: 16.
"""
def __init__(self, in_channels=768, out_channels=[96, 192, 384, 768], readout_type="ignore", patch_size=16):
super(ReassembleBlocks, self).__init__()
assert readout_type in ["ignore", "add", "project"]
self.readout_type = readout_type
self.patch_size = patch_size
self.projects = nn.ModuleList(
[
ConvModule(
in_channels=in_channels,
out_channels=out_channel,
kernel_size=1,
act_layer=None,
)
for out_channel in out_channels
]
)
self.resize_layers = nn.ModuleList(
[
nn.ConvTranspose2d(
in_channels=out_channels[0], out_channels=out_channels[0], kernel_size=4, stride=4, padding=0
),
nn.ConvTranspose2d(
in_channels=out_channels[1], out_channels=out_channels[1], kernel_size=2, stride=2, padding=0
),
nn.Identity(),
nn.Conv2d(
in_channels=out_channels[3], out_channels=out_channels[3], kernel_size=3, stride=2, padding=1
),
]
)
if self.readout_type == "project":
self.readout_projects = nn.ModuleList()
for _ in range(len(self.projects)):
self.readout_projects.append(nn.Sequential(nn.Linear(2 * in_channels, in_channels), nn.GELU()))
def forward(self, inputs):
assert isinstance(inputs, list)
out = []
for i, x in enumerate(inputs):
assert len(x) == 2
x, cls_token = x[0], x[1]
feature_shape = x.shape
if self.readout_type == "project":
x = x.flatten(2).permute((0, 2, 1))
readout = cls_token.unsqueeze(1).expand_as(x)
x = self.readout_projects[i](torch.cat((x, readout), -1))
x = x.permute(0, 2, 1).reshape(feature_shape)
elif self.readout_type == "add":
x = x.flatten(2) + cls_token.unsqueeze(-1)
x = x.reshape(feature_shape)
else:
pass
x = self.projects[i](x)
x = self.resize_layers[i](x)
out.append(x)
return out
class PreActResidualConvUnit(nn.Module):
"""ResidualConvUnit, pre-activate residual unit.
Args:
in_channels (int): number of channels in the input feature map.
act_layer (nn.Module): activation layer.
norm_layer (nn.Module): norm layer.
stride (int): stride of the first block. Default: 1
dilation (int): dilation rate for convs layers. Default: 1.
"""
def __init__(self, in_channels, act_layer, norm_layer, stride=1, dilation=1):
super(PreActResidualConvUnit, self).__init__()
self.conv1 = ConvModule(
in_channels,
in_channels,
3,
stride=stride,
padding=dilation,
dilation=dilation,
norm_layer=norm_layer,
act_layer=act_layer,
bias=False,
order=("act", "conv", "norm"),
)
self.conv2 = ConvModule(
in_channels,
in_channels,
3,
padding=1,
norm_layer=norm_layer,
act_layer=act_layer,
bias=False,
order=("act", "conv", "norm"),
)
def forward(self, inputs):
inputs_ = inputs.clone()
x = self.conv1(inputs)
x = self.conv2(x)
return x + inputs_
class FeatureFusionBlock(nn.Module):
"""FeatureFusionBlock, merge feature map from different stages.
Args:
in_channels (int): Input channels.
act_layer (nn.Module): activation layer for ResidualConvUnit.
norm_layer (nn.Module): normalization layer.
expand (bool): Whether expand the channels in post process block.
Default: False.
align_corners (bool): align_corner setting for bilinear upsample.
Default: True.
"""
def __init__(self, in_channels, act_layer, norm_layer, expand=False, align_corners=True):
super(FeatureFusionBlock, self).__init__()
self.in_channels = in_channels
self.expand = expand
self.align_corners = align_corners
self.out_channels = in_channels
if self.expand:
self.out_channels = in_channels // 2
self.project = ConvModule(self.in_channels, self.out_channels, kernel_size=1, act_layer=None, bias=True)
self.res_conv_unit1 = PreActResidualConvUnit(
in_channels=self.in_channels, act_layer=act_layer, norm_layer=norm_layer
)
self.res_conv_unit2 = PreActResidualConvUnit(
in_channels=self.in_channels, act_layer=act_layer, norm_layer=norm_layer
)
def forward(self, *inputs):
x = inputs[0]
if len(inputs) == 2:
if x.shape != inputs[1].shape:
res = resize(inputs[1], size=(x.shape[2], x.shape[3]), mode="bilinear", align_corners=False)
else:
res = inputs[1]
x = x + self.res_conv_unit1(res)
x = self.res_conv_unit2(x)
x = resize(x, scale_factor=2, mode="bilinear", align_corners=self.align_corners)
x = self.project(x)
return x
class DPTHead(DepthBaseDecodeHead):
"""Vision Transformers for Dense Prediction.
This head is implemented of `DPT <https://arxiv.org/abs/2103.13413>`_.
Args:
embed_dims (int): The embed dimension of the ViT backbone.
Default: 768.
post_process_channels (List): Out channels of post process conv
layers. Default: [96, 192, 384, 768].
readout_type (str): Type of readout operation. Default: 'ignore'.
patch_size (int): The patch size. Default: 16.
expand_channels (bool): Whether expand the channels in post process
block. Default: False.
"""
def __init__(
self,
embed_dims=768,
post_process_channels=[96, 192, 384, 768],
readout_type="ignore",
patch_size=16,
expand_channels=False,
**kwargs,
):
super(DPTHead, self).__init__(**kwargs)
self.in_channels = self.in_channels
self.expand_channels = expand_channels
self.reassemble_blocks = ReassembleBlocks(embed_dims, post_process_channels, readout_type, patch_size)
self.post_process_channels = [
channel * math.pow(2, i) if expand_channels else channel for i, channel in enumerate(post_process_channels)
]
self.convs = nn.ModuleList()
for channel in self.post_process_channels:
self.convs.append(ConvModule(channel, self.channels, kernel_size=3, padding=1, act_layer=None, bias=False))
self.fusion_blocks = nn.ModuleList()
for _ in range(len(self.convs)):
self.fusion_blocks.append(FeatureFusionBlock(self.channels, self.act_layer, self.norm_layer))
self.fusion_blocks[0].res_conv_unit1 = None
self.project = ConvModule(self.channels, self.channels, kernel_size=3, padding=1, norm_layer=self.norm_layer)
self.num_fusion_blocks = len(self.fusion_blocks)
self.num_reassemble_blocks = len(self.reassemble_blocks.resize_layers)
self.num_post_process_channels = len(self.post_process_channels)
assert self.num_fusion_blocks == self.num_reassemble_blocks
assert self.num_reassemble_blocks == self.num_post_process_channels
self.conv_depth = HeadDepth(self.channels)
def forward(self, inputs, img_metas):
assert len(inputs) == self.num_reassemble_blocks
x = [inp for inp in inputs]
x = self.reassemble_blocks(x)
x = [self.convs[i](feature) for i, feature in enumerate(x)]
out = self.fusion_blocks[0](x[-1])
for i in range(1, len(self.fusion_blocks)):
out = self.fusion_blocks[i](out, x[-(i + 1)])
out = self.project(out)
out = self.depth_pred(out)
return out
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
import warnings
import torch.nn.functional as F
def resize(input, size=None, scale_factor=None, mode="nearest", align_corners=None, warning=False):
if warning:
if size is not None and align_corners:
input_h, input_w = tuple(int(x) for x in input.shape[2:])
output_h, output_w = tuple(int(x) for x in size)
if output_h > input_h or output_w > output_h:
if (
(output_h > 1 and output_w > 1 and input_h > 1 and input_w > 1)
and (output_h - 1) % (input_h - 1)
and (output_w - 1) % (input_w - 1)
):
warnings.warn(
f"When align_corners={align_corners}, "
"the output would more aligned if "
f"input size {(input_h, input_w)} is `x+1` and "
f"out size {(output_h, output_w)} is `nx+1`"
)
return F.interpolate(input, size, scale_factor, mode, align_corners)
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.
from .decode_heads import BNHead, DPTHead
from .encoder_decoder import DepthEncoderDecoder
# ORIGINAL LICENSE
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Modified by Zexin He in 2023-2024.
# The modifications are subject to the same license as the original.
import itertools
import torch
import torch.nn as nn
from .utils.renderer import ImportanceRenderer
from .utils.ray_sampler import RaySampler
class ShiftedSoftplus(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return nn.functional.softplus(x - 1)
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)
@torch.compile
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': False,
'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, anchors, resolutions, bg_colors, region_size: int):
# planes: (N, 3, D', H', W')
# cameras: (N, M, D_cam)
# anchors: (N, M, 2)
# resolutions: (N, M, 1)
# bg_colors: (N, M, 1)
# region_size: int
assert planes.shape[0] == cameras.shape[0], "Batch size mismatch for planes and cameras"
assert planes.shape[0] == anchors.shape[0], "Batch size mismatch for planes and anchors"
assert cameras.shape[1] == anchors.shape[1], "Number of views mismatch for cameras and anchors"
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),
resolutions=resolutions.reshape(-1, 1),
anchors=anchors.reshape(-1, 2),
region_size=region_size,
)
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,
bg_colors=bg_colors.reshape(-1, 1),
)
# Reshape into 'raw' neural-rendered image
Himg = Wimg = region_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()
}
return point_features
# 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.
#
# Empty
# 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 in 2023-2024.
# 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
@torch.compile
def forward(self, cam2world_matrix, intrinsics, resolutions, anchors, region_size):
"""
Create batches of rays and return origins and directions.
cam2world_matrix: (N, 4, 4)
intrinsics: (N, 3, 3)
resolutions: (N, 1)
anchors: (N, 2)
region_size: int
ray_origins: (N, M, 3)
ray_dirs: (N, M, 2)
"""
N, M = cam2world_matrix.shape[0], region_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(region_size, dtype=torch.float32, device=cam2world_matrix.device),
torch.arange(region_size, dtype=torch.float32, device=cam2world_matrix.device),
indexing='ij',
))
uv = uv.flip(0).reshape(2, -1).transpose(1, 0)
uv = uv.unsqueeze(0).repeat(cam2world_matrix.shape[0], 1, 1)
# anchors are indexed as normal (row, col) but uv is indexed as (x, y)
x_cam = (uv[:, :, 0].view(N, -1) + anchors[:, 1].unsqueeze(-1)) * (1./resolutions) + (0.5/resolutions)
y_cam = (uv[:, :, 1].view(N, -1) + anchors[:, 0].unsqueeze(-1)) * (1./resolutions) + (0.5/resolutions)
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)
_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)
cam2world_matrix = torch.bmm(cam2world_matrix, _opencv2blender)
world_rel_points = torch.bmm(cam2world_matrix, 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
# 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 in 2023-2024.
# 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, bg_colors=None):
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
else:
assert bg_colors is not None, "Must provide bg_colors if white_back is False"
composite_rgb = composite_rgb + bg_colors.unsqueeze(-1) * (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, bg_colors=None):
composite_rgb, composite_depth, weights = self.run_forward(colors, densities, depths, rendering_options, bg_colors=bg_colors)
return composite_rgb, composite_depth, weights
# 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
# 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 in 2023-2024.
# 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)
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
projected_coordinates = project_onto_planes(plane_axes, coordinates).unsqueeze(1)
output_features = torch.nn.functional.grid_sample(plane_features, 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 = depths.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 = 8
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, bg_colors=None):
# 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'])
# 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, bg_colors=bg_colors)
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, bg_colors=bg_colors)
else:
rgb_final, depth_final, weights = self.ray_marcher(colors_coarse, densities_coarse, depths_coarse, rendering_options, bg_colors=bg_colors)
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
# 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.
# 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 os
from typing import Union
import random
import numpy as np
import torch
from megfile import smart_path_join, smart_open
from .base import BaseDataset
from .cam_utils import build_camera_standard, build_camera_principle, camera_normalization_objaverse
from ..utils.proxy import no_proxy
__all__ = ['ObjaverseDataset']
class ObjaverseDataset(BaseDataset):
def __init__(self, root_dirs: list[str], meta_path: str,
sample_side_views: int,
render_image_res_low: int, render_image_res_high: int, render_region_size: int,
source_image_res: int, normalize_camera: bool,
normed_dist_to_center: Union[float, str] = None, num_all_views: int = 32):
super().__init__(root_dirs, meta_path)
self.sample_side_views = sample_side_views
self.render_image_res_low = render_image_res_low
self.render_image_res_high = render_image_res_high
self.render_region_size = render_region_size
self.source_image_res = source_image_res
self.normalize_camera = normalize_camera
self.normed_dist_to_center = normed_dist_to_center
self.num_all_views = num_all_views
@staticmethod
def _load_pose(file_path):
pose = np.load(smart_open(file_path, 'rb'))
pose = torch.from_numpy(pose).float()
return pose
@no_proxy
def inner_get_item(self, idx):
"""
Loaded contents:
rgbs: [M, 3, H, W]
poses: [M, 3, 4], [R|t]
intrinsics: [3, 2], [[fx, fy], [cx, cy], [weight, height]]
"""
uid = self.uids[idx]
root_dir = self._locate_datadir(self.root_dirs, uid, locator="intrinsics.npy")
pose_dir = os.path.join(root_dir, uid, 'pose')
rgba_dir = os.path.join(root_dir, uid, 'rgba')
intrinsics_path = os.path.join(root_dir, uid, 'intrinsics.npy')
# load intrinsics
intrinsics = np.load(smart_open(intrinsics_path, 'rb'))
intrinsics = torch.from_numpy(intrinsics).float()
# sample views (incl. source view and side views)
sample_views = np.random.choice(range(self.num_all_views), self.sample_side_views + 1, replace=False)
poses, rgbs, bg_colors = [], [], []
source_image = None
for view in sample_views:
pose_path = smart_path_join(pose_dir, f'{view:03d}.npy')
rgba_path = smart_path_join(rgba_dir, f'{view:03d}.png')
pose = self._load_pose(pose_path)
bg_color = random.choice([0.0, 0.5, 1.0])
rgb = self._load_rgba_image(rgba_path, bg_color=bg_color)
poses.append(pose)
rgbs.append(rgb)
bg_colors.append(bg_color)
if source_image is None:
source_image = self._load_rgba_image(rgba_path, bg_color=1.0)
assert source_image is not None, "Really bad luck!"
poses = torch.stack(poses, dim=0)
rgbs = torch.cat(rgbs, dim=0)
if self.normalize_camera:
poses = camera_normalization_objaverse(self.normed_dist_to_center, poses)
# build source and target camera features
source_camera = build_camera_principle(poses[:1], intrinsics.unsqueeze(0)).squeeze(0)
render_camera = build_camera_standard(poses, intrinsics.repeat(poses.shape[0], 1, 1))
# adjust source image resolution
source_image = torch.nn.functional.interpolate(
source_image, size=(self.source_image_res, self.source_image_res), mode='bicubic', align_corners=True).squeeze(0)
source_image = torch.clamp(source_image, 0, 1)
# adjust render image resolution and sample intended rendering region
render_image_res = np.random.randint(self.render_image_res_low, self.render_image_res_high + 1)
render_image = torch.nn.functional.interpolate(
rgbs, size=(render_image_res, render_image_res), mode='bicubic', align_corners=True)
render_image = torch.clamp(render_image, 0, 1)
anchors = torch.randint(
0, render_image_res - self.render_region_size + 1, size=(self.sample_side_views + 1, 2))
crop_indices = torch.arange(0, self.render_region_size, device=render_image.device)
index_i = (anchors[:, 0].unsqueeze(1) + crop_indices).view(-1, self.render_region_size, 1)
index_j = (anchors[:, 1].unsqueeze(1) + crop_indices).view(-1, 1, self.render_region_size)
batch_indices = torch.arange(self.sample_side_views + 1, device=render_image.device).view(-1, 1, 1)
cropped_render_image = render_image[batch_indices, :, index_i, index_j].permute(0, 3, 1, 2)
return {
'uid': uid,
'source_camera': source_camera,
'render_camera': render_camera,
'source_image': source_image,
'render_image': cropped_render_image,
'render_anchors': anchors,
'render_full_resolutions': torch.tensor([[render_image_res]], dtype=torch.float32).repeat(self.sample_side_views + 1, 1),
'render_bg_colors': torch.tensor(bg_colors, dtype=torch.float32).unsqueeze(-1),
}
# 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 math
from functools import partial
import torch
__all__ = ['MixerDataset']
class MixerDataset(torch.utils.data.Dataset):
def __init__(self,
split: str,
subsets: list[dict],
**dataset_kwargs,
):
self.subsets = [
self._dataset_fn(subset, split)(**dataset_kwargs)
for subset in subsets
]
self.virtual_lens = [
math.ceil(subset_config['sample_rate'] * len(subset_obj))
for subset_config, subset_obj in zip(subsets, self.subsets)
]
@staticmethod
def _dataset_fn(subset_config: dict, split: str):
name = subset_config['name']
dataset_cls = None
if name == "objaverse":
from .objaverse import ObjaverseDataset
dataset_cls = ObjaverseDataset
# elif name == 'mvimgnet':
# from .mvimgnet import MVImgNetDataset
# dataset_cls = MVImgNetDataset
else:
raise NotImplementedError(f"Dataset {name} not implemented")
return partial(
dataset_cls,
root_dirs=subset_config['root_dirs'],
meta_path=subset_config['meta_path'][split],
)
def __len__(self):
return sum(self.virtual_lens)
def __getitem__(self, idx):
subset_idx = 0
virtual_idx = idx
while virtual_idx >= self.virtual_lens[subset_idx]:
virtual_idx -= self.virtual_lens[subset_idx]
subset_idx += 1
real_idx = virtual_idx % len(self.subsets[subset_idx])
return self.subsets[subset_idx][real_idx]
# 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 math
import torch
"""
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=RT.dtype, device=RT.device).repeat(RT.shape[0], 1, 1)
], 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 camera_normalization_objaverse(normed_dist_to_center, poses: torch.Tensor, ret_transform: bool = False):
assert normed_dist_to_center is not None
pivotal_pose = compose_extrinsic_RT(poses[:1])
dist_to_center = pivotal_pose[:, :3, 3].norm(dim=-1, keepdim=True).item() \
if normed_dist_to_center == 'auto' else normed_dist_to_center
# compute camera norm (new version)
canonical_camera_extrinsics = torch.tensor([[
[1, 0, 0, 0],
[0, 0, -1, -dist_to_center],
[0, 1, 0, 0],
[0, 0, 0, 1],
]], dtype=torch.float32)
pivotal_pose_inv = torch.inverse(pivotal_pose)
camera_norm_matrix = torch.bmm(canonical_camera_extrinsics, pivotal_pose_inv)
# normalize all views
poses = compose_extrinsic_RT(poses)
poses = torch.bmm(camera_norm_matrix.repeat(poses.shape[0], 1, 1), poses)
poses = decompose_extrinsic_RT(poses)
if ret_transform:
return poses, camera_norm_matrix.squeeze(dim=0)
return poses
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,
device: torch.device = torch.device('cpu'),
):
"""
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, device=device)
if up_world is None:
up_world = torch.tensor([0, 0, 1], dtype=torch.float32, device=device)
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 surrounding_views_linspace(n_views: int, radius: float = 2.0, height: float = 0.8, device: torch.device = torch.device('cpu')):
"""
n_views: number of surrounding views
radius: camera dist to center
height: height of the camera
return: (M, 3, 4)
"""
assert n_views > 0
assert radius > 0
theta = torch.linspace(-torch.pi / 2, 3 * torch.pi / 2, n_views, device=device)
projected_radius = math.sqrt(radius ** 2 - height ** 2)
x = torch.cos(theta) * projected_radius
y = torch.sin(theta) * projected_radius
z = torch.full((n_views,), height, device=device)
camera_positions = torch.stack([x, y, z], dim=1)
extrinsics = center_looking_at_camera_pose(camera_positions, device=device)
return extrinsics
def create_intrinsics(
f: float,
c: float = None, cx: float = None, cy: float = None,
w: float = 1., h: float = 1.,
dtype: torch.dtype = torch.float32,
device: torch.device = torch.device('cpu'),
):
"""
return: (3, 2)
"""
fx = fy = f
if c is not None:
assert cx is None and cy is None, "c and cx/cy cannot be used together"
cx = cy = c
else:
assert cx is not None and cy is not None, "cx/cy must be provided when c is not provided"
fx, fy, cx, cy, w, h = fx/w, fy/h, cx/w, cy/h, 1., 1.
intrinsics = torch.tensor([
[fx, fy],
[cx, cy],
[w, h],
], dtype=dtype, device=device)
return intrinsics
# 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.
from .mixer import MixerDataset
# 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.
from abc import ABC, abstractmethod
import json
import numpy as np
import torch
from PIL import Image
from megfile import smart_open, smart_path_join, smart_exists
class BaseDataset(torch.utils.data.Dataset, ABC):
def __init__(self, root_dirs: list[str], meta_path: str):
super().__init__()
self.root_dirs = root_dirs
self.uids = self._load_uids(meta_path)
def __len__(self):
return len(self.uids)
@abstractmethod
def inner_get_item(self, idx):
pass
def __getitem__(self, idx):
try:
return self.inner_get_item(idx)
except Exception as e:
print(f"[DEBUG-DATASET] Error when loading {self.uids[idx]}")
# return self.__getitem__(idx+1)
raise e
@staticmethod
def _load_uids(meta_path: str):
# meta_path is a json file
with open(meta_path, 'r') as f:
uids = json.load(f)
return uids
@staticmethod
def _load_rgba_image(file_path, bg_color: float = 1.0):
''' Load and blend RGBA image to RGB with certain background, 0-1 scaled '''
rgba = np.array(Image.open(smart_open(file_path, 'rb')))
rgba = torch.from_numpy(rgba).float() / 255.0
rgba = rgba.permute(2, 0, 1).unsqueeze(0)
rgb = rgba[:, :3, :, :] * rgba[:, 3:4, :, :] + bg_color * (1 - rgba[:, 3:, :, :])
rgba[:, :3, ...] * rgba[:, 3:, ...] + (1 - rgba[:, 3:, ...])
return rgb
@staticmethod
def _locate_datadir(root_dirs, uid, locator: str):
for root_dir in root_dirs:
datadir = smart_path_join(root_dir, uid, locator)
if smart_exists(datadir):
return root_dir
raise FileNotFoundError(f"Cannot find valid data directory for uid {uid}")
# 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.
from accelerate.logging import get_logger
logger = get_logger(__name__)
def configure_dynamo(config: dict):
try:
import torch._dynamo
logger.debug(f'Configuring torch._dynamo.config with {config}')
for k, v in config.items():
if v is None:
logger.debug(f'Skipping torch._dynamo.config.{k} with None')
continue
if hasattr(torch._dynamo.config, k):
logger.warning(f'Overriding torch._dynamo.config.{k} from {getattr(torch._dynamo.config, k)} to {v}')
setattr(torch._dynamo.config, k, v)
except ImportError:
logger.debug('torch._dynamo not found, skipping')
pass
# 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.
from torch.profiler import profile
class DummyProfiler(profile):
def __init__(self):
pass
def __enter__(self):
return self
def __exit__(self, *args):
pass
def step(self):
pass
# 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 os
NO_PROXY = "OPENLRM_NO_DATA_PROXY" in os.environ
def no_proxy(func):
"""Decorator to disable proxy but then restore after the function call."""
def wrapper(*args, **kwargs):
# http_proxy, https_proxy, HTTP_PROXY, HTTPS_PROXY, all_proxy
http_proxy = os.environ.get('http_proxy')
https_proxy = os.environ.get('https_proxy')
HTTP_PROXY = os.environ.get('HTTP_PROXY')
HTTPS_PROXY = os.environ.get('HTTPS_PROXY')
all_proxy = os.environ.get('all_proxy')
os.environ['http_proxy'] = ''
os.environ['https_proxy'] = ''
os.environ['HTTP_PROXY'] = ''
os.environ['HTTPS_PROXY'] = ''
os.environ['all_proxy'] = ''
try:
return func(*args, **kwargs)
finally:
os.environ['http_proxy'] = http_proxy
os.environ['https_proxy'] = https_proxy
os.environ['HTTP_PROXY'] = HTTP_PROXY
os.environ['HTTPS_PROXY'] = HTTPS_PROXY
os.environ['all_proxy'] = all_proxy
if NO_PROXY:
return wrapper
else:
return func
# 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 math
from torch.optim.lr_scheduler import LRScheduler
from accelerate.logging import get_logger
logger = get_logger(__name__)
class CosineWarmupScheduler(LRScheduler):
def __init__(self, optimizer, warmup_iters: int, max_iters: int, initial_lr: float = 1e-10, last_iter: int = -1):
self.warmup_iters = warmup_iters
self.max_iters = max_iters
self.initial_lr = initial_lr
super().__init__(optimizer, last_iter)
def get_lr(self):
logger.debug(f"step count: {self._step_count} | warmup iters: {self.warmup_iters} | max iters: {self.max_iters}")
if self._step_count <= self.warmup_iters:
return [
self.initial_lr + (base_lr - self.initial_lr) * self._step_count / self.warmup_iters
for base_lr in self.base_lrs]
else:
cos_iter = self._step_count - self.warmup_iters
cos_max_iter = self.max_iters - self.warmup_iters
cos_theta = cos_iter / cos_max_iter * math.pi
cos_lr = [base_lr * (1 + math.cos(cos_theta)) / 2 for base_lr in self.base_lrs]
return cos_lr
# 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.
class Registry:
"""Registry class"""
def __init__(self):
self._registry = {}
def register(self, name):
"""Register a module"""
def decorator(cls):
assert name not in self._registry, 'Module {} already registered'.format(name)
self._registry[name] = cls
return cls
return decorator
def __getitem__(self, name):
"""Get a module"""
return self._registry[name]
def __contains__(self, name):
return name in self._registry
# 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 numpy as np
import rembg
import cv2
class Preprocessor:
"""
Preprocessing under cv2 conventions.
"""
def __init__(self):
self.rembg_session = rembg.new_session(
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)
def preprocess(self, image_path: str, save_path: str, rmbg: bool = True, recenter: bool = True, size: int = 512, border_ratio: float = 0.2):
image = self.step_load_to_size(image_path=image_path, size=size*2)
if rmbg:
image = self.step_rembg(image_in=image)
else:
image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA)
if recenter:
image = self.step_recenter(image_in=image, border_ratio=border_ratio, square_size=size)
else:
image = cv2.resize(
src=image,
dsize=(size, size),
interpolation=cv2.INTER_AREA,
)
return cv2.imwrite(save_path, image)
def step_rembg(self, image_in: np.ndarray) -> np.ndarray:
image_out = rembg.remove(
data=image_in,
session=self.rembg_session,
)
return image_out
def step_recenter(self, image_in: np.ndarray, border_ratio: float, square_size: int) -> np.ndarray:
assert image_in.shape[-1] == 4, "Image to recenter must be RGBA"
mask = image_in[..., -1] > 0
ijs = np.nonzero(mask)
# find bbox
i_min, i_max = ijs[0].min(), ijs[0].max()
j_min, j_max = ijs[1].min(), ijs[1].max()
bbox_height, bbox_width = i_max - i_min, j_max - j_min
# recenter and resize
desired_size = int(square_size * (1 - border_ratio))
scale = desired_size / max(bbox_height, bbox_width)
desired_height, desired_width = int(bbox_height * scale), int(bbox_width * scale)
desired_i_min, desired_j_min = (square_size - desired_height) // 2, (square_size - desired_width) // 2
desired_i_max, desired_j_max = desired_i_min + desired_height, desired_j_min + desired_width
# create new image
image_out = np.zeros((square_size, square_size, 4), dtype=np.uint8)
image_out[desired_i_min:desired_i_max, desired_j_min:desired_j_max] = cv2.resize(
src=image_in[i_min:i_max, j_min:j_max],
dsize=(desired_width, desired_height),
interpolation=cv2.INTER_AREA,
)
return image_out
def step_load_to_size(self, image_path: str, size: int) -> np.ndarray:
image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
height, width = image.shape[:2]
scale = size / max(height, width)
height, width = int(height * scale), int(width * scale)
image_out = cv2.resize(
src=image,
dsize=(width, height),
interpolation=cv2.INTER_AREA,
)
return image_out
# 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 os
import numpy as np
import imageio
def images_to_video(images, output_path, fps, gradio_codec: bool, verbose=False):
# images: (T, C, H, W)
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)
frames = np.stack(frames)
if gradio_codec:
imageio.mimwrite(output_path, frames, fps=fps, quality=10)
else:
imageio.mimwrite(output_path, frames, fps=fps, codec='mpeg4', quality=10)
if verbose:
print(f"Using gradio codec option {gradio_codec}")
print(f"Saved video to {output_path}")
# 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 os
import logging
from tqdm.auto import tqdm
class TqdmStreamHandler(logging.StreamHandler):
def emit(self, record):
tqdm.write(self.format(record))
def configure_logger(stream_level, log_level, file_path = None):
_stream_level = stream_level.upper()
_log_level = log_level.upper()
_project_level = _log_level
_formatter = logging.Formatter("[%(asctime)s] %(name)s: [%(levelname)s] %(message)s")
_stream_handler = TqdmStreamHandler()
_stream_handler.setLevel(_stream_level)
_stream_handler.setFormatter(_formatter)
if file_path is not None:
os.makedirs(os.path.dirname(file_path), exist_ok=True)
_file_handler = logging.FileHandler(file_path)
_file_handler.setLevel(_log_level)
_file_handler.setFormatter(_formatter)
_project_logger = logging.getLogger(__name__.split('.')[0])
_project_logger.setLevel(_project_level)
_project_logger.addHandler(_stream_handler)
if file_path is not None:
_project_logger.addHandler(_file_handler)
# 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.nn as nn
from huggingface_hub import PyTorchModelHubMixin
def wrap_model_hub(model_cls: nn.Module):
class HfModel(model_cls, PyTorchModelHubMixin):
def __init__(self, config: dict):
super().__init__(**config)
self.config = config
return HfModel
# 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.
#
# Empty
# 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 sys
sys.path.append('.')
import argparse
from accelerate import Accelerator
from openlrm.utils.hf_hub import wrap_model_hub
from openlrm.models import model_dict
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--model_type', type=str, required=True)
parser.add_argument('--local_ckpt', type=str, required=True)
parser.add_argument('--repo_id', type=str, required=True)
args, unknown = parser.parse_known_args()
accelerator = Accelerator()
hf_model_cls = wrap_model_hub(model_dict[args.model_type])
hf_model = hf_model_cls.from_pretrained(args.local_ckpt)
hf_model.push_to_hub(
repo_id=args.repo_id,
config=hf_model.config,
private=True,
)
# 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 argparse
from omegaconf import OmegaConf
from megfile import smart_path_join, smart_exists, smart_listdir, smart_makedirs, smart_copy
from tempfile import TemporaryDirectory
import torch.nn as nn
from accelerate import Accelerator
import safetensors
import sys
sys.path.append(".")
from openlrm.utils.hf_hub import wrap_model_hub
from openlrm.utils.proxy import no_proxy
from openlrm.models import model_dict
@no_proxy
def auto_load_model(cfg, model: nn.Module) -> int:
ckpt_root = smart_path_join(
cfg.saver.checkpoint_root,
cfg.experiment.parent, cfg.experiment.child,
)
if not smart_exists(ckpt_root):
raise FileNotFoundError(f"Checkpoint root not found: {ckpt_root}")
ckpt_dirs = smart_listdir(ckpt_root)
if len(ckpt_dirs) == 0:
raise FileNotFoundError(f"No checkpoint found in {ckpt_root}")
ckpt_dirs.sort()
load_step = f"{cfg.convert.global_step}" if cfg.convert.global_step is not None else ckpt_dirs[-1]
load_model_path = smart_path_join(ckpt_root, load_step, 'model.safetensors')
if load_model_path.startswith("s3"):
tmpdir = TemporaryDirectory()
tmp_model_path = smart_path_join(tmpdir.name, f"tmp.safetensors")
smart_copy(load_model_path, tmp_model_path)
load_model_path = tmp_model_path
print(f"Loading from {load_model_path}")
safetensors.torch.load_model(model, load_model_path)
return int(load_step)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default='./assets/config.yaml')
args, unknown = parser.parse_known_args()
cfg = OmegaConf.load(args.config)
cli_cfg = OmegaConf.from_cli(unknown)
cfg = OmegaConf.merge(cfg, cli_cfg)
"""
[cfg.convert]
global_step: int
save_dir: str
"""
accelerator = Accelerator()
hf_model_cls = wrap_model_hub(model_dict[cfg.experiment.type])
hf_model = hf_model_cls(OmegaConf.to_container(cfg.model))
loaded_step = auto_load_model(cfg, hf_model)
dump_path = smart_path_join(
f"./exps/releases",
cfg.experiment.parent, cfg.experiment.child,
f'step_{loaded_step:06d}',
)
print(f"Saving locally to {dump_path}")
smart_makedirs(dump_path, exist_ok=True)
hf_model.save_pretrained(
save_directory=dump_path,
config=hf_model.config,
)
"""
Blender script to render images of 3D models.
This script is designed to render data used in the [OpenLRM project](https://github.com/3DTopia/OpenLRM).
Modified from https://github.com/cvlab-columbia/zero123/blob/main/objaverse-rendering/scripts/blender_script.py
Original script licensed under MIT, found at the root of its repository.
Modifications are licensed under Apache 2.0.
"""
import argparse
import math
import os
import random
import sys
import time
import urllib.request
from typing import Tuple
from mathutils import Vector
import numpy as np
import bpy
parser = argparse.ArgumentParser()
parser.add_argument(
"--object_path",
type=str,
required=True,
help="Path to the object file",
)
parser.add_argument("--output_dir", type=str, default="./views")
parser.add_argument(
"--engine", type=str, default="CYCLES", choices=["CYCLES", "BLENDER_EEVEE"]
)
parser.add_argument("--num_images", type=int, default=32)
parser.add_argument("--resolution", type=int, default=1024)
argv = sys.argv[sys.argv.index("--") + 1 :]
args = parser.parse_args(argv)
print('===================', args.engine, '===================')
context = bpy.context
scene = context.scene
render = scene.render
render.engine = args.engine
render.image_settings.file_format = "PNG"
render.image_settings.color_mode = "RGBA"
render.resolution_x = args.resolution
render.resolution_y = args.resolution
render.resolution_percentage = 100
scene.cycles.device = "GPU"
scene.cycles.samples = 128
scene.cycles.diffuse_bounces = 1
scene.cycles.glossy_bounces = 1
scene.cycles.transparent_max_bounces = 3
scene.cycles.transmission_bounces = 3
scene.cycles.filter_width = 0.01
scene.cycles.use_denoising = True
scene.render.film_transparent = True
# Set the device_type
cycles_preferences = bpy.context.preferences.addons["cycles"].preferences
cycles_preferences.compute_device_type = "CUDA" # or "OPENCL"
cuda_devices = cycles_preferences.get_devices_for_type("CUDA")
for device in cuda_devices:
device.use = True
def compose_RT(R, T):
return np.hstack((R, T.reshape(-1, 1)))
def sample_point_on_sphere(radius: float) -> Tuple[float, float, float]:
theta = random.random() * 2 * math.pi
phi = math.acos(2 * random.random() - 1)
return (
radius * math.sin(phi) * math.cos(theta),
radius * math.sin(phi) * math.sin(theta),
radius * math.cos(phi),
)
def sample_spherical(radius=3.0, maxz=3.0, minz=0.):
correct = False
while not correct:
vec = np.random.uniform(-1, 1, 3)
vec[2] = np.abs(vec[2])
vec = vec / np.linalg.norm(vec, axis=0) * radius
if maxz > vec[2] > minz:
correct = True
return vec
def sample_spherical(radius_min=1.5, radius_max=2.0, maxz=1.6, minz=-0.75):
correct = False
while not correct:
vec = np.random.uniform(-1, 1, 3)
# vec[2] = np.abs(vec[2])
radius = np.random.uniform(radius_min, radius_max, 1)
vec = vec / np.linalg.norm(vec, axis=0) * radius[0]
if maxz > vec[2] > minz:
correct = True
return vec
def set_camera_location(camera, option: str):
assert option in ['fixed', 'random', 'front']
if option == 'fixed':
x, y, z = 0, -2.25, 0
elif option == 'random':
# from https://blender.stackexchange.com/questions/18530/
x, y, z = sample_spherical(radius_min=1.9, radius_max=2.6, maxz=1.60, minz=-0.75)
elif option == 'front':
x, y, z = 0, -np.random.uniform(1.9, 2.6, 1)[0], 0
camera.location = x, y, z
# adjust orientation
direction = - camera.location
rot_quat = direction.to_track_quat('-Z', 'Y')
camera.rotation_euler = rot_quat.to_euler()
return camera
def add_lighting(option: str) -> None:
assert option in ['fixed', 'random']
# delete the default light
bpy.data.objects["Light"].select_set(True)
bpy.ops.object.delete()
# add a new light
bpy.ops.object.light_add(type="AREA")
light = bpy.data.lights["Area"]
if option == 'fixed':
light.energy = 30000
bpy.data.objects["Area"].location[0] = 0
bpy.data.objects["Area"].location[1] = 1
bpy.data.objects["Area"].location[2] = 0.5
elif option == 'random':
light.energy = random.uniform(80000, 120000)
bpy.data.objects["Area"].location[0] = random.uniform(-2., 2.)
bpy.data.objects["Area"].location[1] = random.uniform(-2., 2.)
bpy.data.objects["Area"].location[2] = random.uniform(1.0, 3.0)
# set light scale
bpy.data.objects["Area"].scale[0] = 200
bpy.data.objects["Area"].scale[1] = 200
bpy.data.objects["Area"].scale[2] = 200
def reset_scene() -> None:
"""Resets the scene to a clean state."""
# delete everything that isn't part of a camera or a light
for obj in bpy.data.objects:
if obj.type not in {"CAMERA", "LIGHT"}:
bpy.data.objects.remove(obj, do_unlink=True)
# delete all the materials
for material in bpy.data.materials:
bpy.data.materials.remove(material, do_unlink=True)
# delete all the textures
for texture in bpy.data.textures:
bpy.data.textures.remove(texture, do_unlink=True)
# delete all the images
for image in bpy.data.images:
bpy.data.images.remove(image, do_unlink=True)
# load the glb model
def load_object(object_path: str) -> None:
"""Loads a glb model into the scene."""
if object_path.endswith(".glb"):
bpy.ops.import_scene.gltf(filepath=object_path, merge_vertices=True)
elif object_path.endswith(".fbx"):
bpy.ops.import_scene.fbx(filepath=object_path)
else:
raise ValueError(f"Unsupported file type: {object_path}")
def scene_bbox(single_obj=None, ignore_matrix=False):
bbox_min = (math.inf,) * 3
bbox_max = (-math.inf,) * 3
found = False
for obj in scene_meshes() if single_obj is None else [single_obj]:
found = True
for coord in obj.bound_box:
coord = Vector(coord)
if not ignore_matrix:
coord = obj.matrix_world @ coord
bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
if not found:
raise RuntimeError("no objects in scene to compute bounding box for")
return Vector(bbox_min), Vector(bbox_max)
def scene_root_objects():
for obj in bpy.context.scene.objects.values():
if not obj.parent:
yield obj
def scene_meshes():
for obj in bpy.context.scene.objects.values():
if isinstance(obj.data, (bpy.types.Mesh)):
yield obj
def normalize_scene(box_scale: float):
bbox_min, bbox_max = scene_bbox()
scale = box_scale / max(bbox_max - bbox_min)
for obj in scene_root_objects():
obj.scale = obj.scale * scale
# Apply scale to matrix_world.
bpy.context.view_layer.update()
bbox_min, bbox_max = scene_bbox()
offset = -(bbox_min + bbox_max) / 2
for obj in scene_root_objects():
obj.matrix_world.translation += offset
bpy.ops.object.select_all(action="DESELECT")
def setup_camera():
cam = scene.objects["Camera"]
cam.location = (0, 1.2, 0)
cam.data.lens = 24
cam.data.sensor_width = 32
cam.data.sensor_height = 32 # affects instrinsics calculation, should be set explicitly
cam_constraint = cam.constraints.new(type="TRACK_TO")
cam_constraint.track_axis = "TRACK_NEGATIVE_Z"
cam_constraint.up_axis = "UP_Y"
return cam, cam_constraint
def save_images(object_file: str) -> None:
"""Saves rendered images of the object in the scene."""
os.makedirs(args.output_dir, exist_ok=True)
reset_scene()
# load the object
load_object(object_file)
object_uid = os.path.basename(object_file).split(".")[0]
normalize_scene(box_scale=2)
add_lighting(option='random')
camera, cam_constraint = setup_camera()
# create an empty object to track
empty = bpy.data.objects.new("Empty", None)
scene.collection.objects.link(empty)
cam_constraint.target = empty
# prepare to save
img_dir = os.path.join(args.output_dir, object_uid, 'rgba')
pose_dir = os.path.join(args.output_dir, object_uid, 'pose')
os.makedirs(img_dir, exist_ok=True)
os.makedirs(pose_dir, exist_ok=True)
for i in range(args.num_images):
# set the camera position
camera_option = 'random' if i > 0 else 'front'
camera = set_camera_location(camera, option=camera_option)
# render the image
render_path = os.path.join(img_dir, f"{i:03d}.png")
scene.render.filepath = render_path
bpy.ops.render.render(write_still=True)
# save camera RT matrix (C2W)
location, rotation = camera.matrix_world.decompose()[0:2]
RT = compose_RT(rotation.to_matrix(), np.array(location))
RT_path = os.path.join(pose_dir, f"{i:03d}.npy")
np.save(RT_path, RT)
# save the camera intrinsics
intrinsics = get_calibration_matrix_K_from_blender(camera.data, return_principles=True)
with open(os.path.join(args.output_dir, object_uid,'intrinsics.npy'), 'wb') as f_intrinsics:
np.save(f_intrinsics, intrinsics)
def download_object(object_url: str) -> str:
"""Download the object and return the path."""
# uid = uuid.uuid4()
uid = object_url.split("/")[-1].split(".")[0]
tmp_local_path = os.path.join("tmp-objects", f"{uid}.glb" + ".tmp")
local_path = os.path.join("tmp-objects", f"{uid}.glb")
# wget the file and put it in local_path
os.makedirs(os.path.dirname(tmp_local_path), exist_ok=True)
urllib.request.urlretrieve(object_url, tmp_local_path)
os.rename(tmp_local_path, local_path)
# get the absolute path
local_path = os.path.abspath(local_path)
return local_path
def get_calibration_matrix_K_from_blender(camera, return_principles=False):
"""
Get the camera intrinsic matrix from Blender camera.
Return also numpy array of principle parameters if specified.
Intrinsic matrix K has the following structure in pixels:
[fx 0 cx]
[0 fy cy]
[0 0 1]
Specified principle parameters are:
[fx, fy] - focal lengths in pixels
[cx, cy] - optical centers in pixels
[width, height] - image resolution in pixels
"""
# Render resolution
render = bpy.context.scene.render
width = render.resolution_x * render.pixel_aspect_x
height = render.resolution_y * render.pixel_aspect_y
# Camera parameters
focal_length = camera.lens # Focal length in millimeters
sensor_width = camera.sensor_width # Sensor width in millimeters
sensor_height = camera.sensor_height # Sensor height in millimeters
# Calculate the focal length in pixel units
focal_length_x = width * (focal_length / sensor_width)
focal_length_y = height * (focal_length / sensor_height)
# Assuming the optical center is at the center of the sensor
optical_center_x = width / 2
optical_center_y = height / 2
# Constructing the intrinsic matrix
K = np.array([[focal_length_x, 0, optical_center_x],
[0, focal_length_y, optical_center_y],
[0, 0, 1]])
if return_principles:
return np.array([
[focal_length_x, focal_length_y],
[optical_center_x, optical_center_y],
[width, height],
])
else:
return K
if __name__ == "__main__":
try:
start_i = time.time()
if args.object_path.startswith("http"):
local_path = download_object(args.object_path)
else:
local_path = args.object_path
save_images(local_path)
end_i = time.time()
print("Finished", local_path, "in", end_i - start_i, "seconds")
# delete the object if it was downloaded
if args.object_path.startswith("http"):
os.remove(local_path)
except Exception as e:
print("Failed to render", args.object_path)
print(e)
LRM: LARGE RECONSTRUCTION MODEL FOR
SINGLE IMAGE TO 3D
Yicong Hong12˚ Kai Zhang1
Jiuxiang Gu1
Sai Bi1
Yang Zhou1
Difan Liu1
Feng Liu1
Kalyan Sunkavalli1
Trung Bui1
Hao Tan1
1Adobe Research
2Australian National Univeristy
mr.yiconghong@gmail.com
{kaiz,jigu,sbi,yazhou,diliu,fengl,sunkaval,bui,hatan}@adobe.com
ABSTRACT
We propose the first Large Reconstruction Model (LRM) that predicts the 3D
model of an object from a single input image within just 5 seconds. In contrast to
many previous methods that are trained on small-scale datasets such as ShapeNet
in a category-specific fashion, LRM adopts a highly scalable transformer-based
architecture with 500 million learnable parameters to directly predict a neural ra-
diance field (NeRF) from the input image. We train our model in an end-to-end
manner on massive multi-view data containing around 1 million objects, includ-
ing both synthetic renderings from Objaverse and real captures from MVImgNet.
This combination of a high-capacity model and large-scale training data empowers
our model to be highly generalizable and produce high-quality 3D reconstructions
from various testing inputs, including real-world in-the-wild captures and images
created by generative models. Video demos and interactable 3D meshes can be
found on our LRM project webpage: https://yiconghong.me/LRM.
1
INTRODUCTION
Imagine if we could instantly create a 3D shape from a single image of an arbitrary object. Broad
applications in industrial design, animation, gaming, and AR/VR have strongly motivated relevant
research in seeking a generic and efficient approach towards this long-standing goal. Due to the
underlying ambiguity of 3D geometry in a single view, early learning-based methods usually per-
form well on specific categories, utilizing the category data prior to infer the overall shape (Yu et al.,
2021). Recently, advances in image generation, such as DALL-E (Ramesh et al., 2021) and Stable
Diffusion (Rombach et al., 2022), have inspired research that leverages the remarkable generaliza-
tion capability of 2D diffusion models to enable multi-view supervision (Liu et al., 2023b; Tang
et al., 2023). However, many of these methods require delicate parameter tuning and regularization,
and their results are limited by the pre-trained 2D generative models. Meanwhile, there are many
approaches that rely on per-shape optimization (e.g. optimize a NeRF (Mildenhall et al., 2021;
Chan et al., 2022; Chen et al., 2022a; M¨
uller et al., 2022; Sun et al., 2022)) to construct a consistent
geometry; this process is often slow and impractical.
On the other hand, the great success in natural language processing (Devlin et al., 2018; Brown
et al., 2020; Chowdhery et al., 2022) and image processing (Caron et al., 2021; Radford et al.,
2021; Alayrac et al., 2022; Ramesh et al., 2022) can be largely credited to three critical factors:
(1) using highly scalable and effective neural networks, such as the Transformers (Vaswani et al.,
2017), for modeling the data distribution, (2) enormous datasets for learning generic priors, as well
as (3) self-supervised-like training objectives that encourage the model to discover the underlying
data structure while maintaining high scalability. For instance, the GPT (generative pre-trained
transformer) series (Radford et al., 2019; Brown et al., 2020; OpenAI, 2023) build large language
models with huge transformer networks, large-scale data, and the simple next-word prediction task.
In light of this, we pose the same question for 3D: given sufficient 3D data and a large-scale training
framework, is it possible to learn a generic 3D prior for reconstructing an object from a single
image?
˚Intern at Adobe Research.
1
In this paper, we propose a Large Reconstruction Model (LRM) for single-image to 3D. Our method
adopts a large transformer-based encoder-decoder architecture for learning 3D representations of ob-
jects from a single image in a data-driven manner. Our method takes an image as input and regresses
a NeRF in the form of a triplane representation (Chan et al., 2022). Specifically, LRM utilizes the
pre-trained visual transformer DINO (Caron et al., 2021) as the image encoder to generate the image
features, and learns an image-to-triplane transformer decoder to project the 2D image features onto
the 3D triplane via cross-attention and model the relations among the spatially-structured triplane
tokens via self-attention. The output tokens from the decoder are reshaped and upsampled to the
final triplane feature maps. Afterwards, we can render the images at an arbitrary view by decoding
the triplane feature of each point with an additional shared multi-layer perception (MLP) to get its
color and density and performing volume rendering.
The overall design of LRM maintains high scalability and efficiency. In addition to the use of a fully
transformer-based pipeline, a triplane NeRF is a concise and scalable 3D representation since it is
computationally friendly compared to other representations such as volumes and point clouds. It
also has a better locality with respect to the image input compared to tokenizing the NeRF’s model
weights as in Shap-E (Jun & Nichol, 2023). Moreover, our LRM is trained by simply minimizing the
difference between the rendered images and ground truth images at novel views, without excessive
3D-aware regularization or delicate hyper-parameter tuning, allowing the model to be very efficient
in training and adaptable to a wide range of multi-view image datasets.
To the best of our knowledge, LRM is the first large-scale 3D reconstruction model; it contains more
than 500 million learnable parameters, and it is trained on approximately one million 3D shapes and
video data across diverse categories (Deitke et al., 2023; Yu et al., 2023); this is substantially larger
than recent methods that apply relatively shallower networks and smaller datasets (Chang et al.,
2015; Reizenstein et al., 2021; Downs et al., 2022). Through experiments, we show that LRM can
reconstruct high-fidelity 3D shapes from a wide range of images captured in the real world, as well
as images created by generative models. LRM is also a highly practical solution for downstream
applications since it can produce a 3D shape in just five seconds1 without post-optimization.
2
RELATED WORK
Single Image to 3D Reconstruction
Extensive efforts have been devoted to address this prob-
lem, including early learning-based methods that explore point clouds (Fan et al., 2017; Wu et al.,
2020), voxels (Choy et al., 2016; Tulsiani et al., 2017; Chen & Zhang, 2019), and meshes (Wang
et al., 2018; Gkioxari et al., 2019), as well as various approaches that learn implicit representa-
tions such as SDFs (Park et al., 2019; Mittal et al., 2022), occupancy networks (Mescheder et al.,
2019), and NeRF (Jang & Agapito, 2021; M¨
uller et al., 2022). Leveraging 3D templates (Roth et al.,
2016; Goel et al., 2020; Kanazawa et al., 2018; Kulkarni et al., 2020), semantics (Li et al., 2020),
and poses (Bogo et al., 2016; Novotny et al., 2019) as shape priors have also been widely studied
in category-specific reconstruction. Category-agnostic methods show great generalization poten-
tial (Yan et al., 2016; Niemeyer et al., 2020), but they often unable to produce fine-grained details
even when exploiting spatially-aligned local image features (Xu et al., 2019; Yu et al., 2021).
Very recently, there is an emerging trend of using pre-trained image/language models (Radford et al.,
2021; Li et al., 2022; 2023b; Saharia et al., 2022; Rombach et al., 2022), to introduce semantics and
multi-view guidance for image-to-3D reconstruction (Liu et al., 2023b; Tang et al., 2023; Deng et al.,
2023; Shen et al., 2023b; Anciukeviˇ
cius et al., 2023; Melas-Kyriazi et al., 2023; Metzer et al., 2023;
Xu et al., 2023; Qian et al., 2023; Li et al., 2023a). For instance, Zero-1-to-3 fine-tunes the Stable
Diffusion model to generate novel views by conditioning on the input image and camera poses (Liu
et al., 2023b); its view consistency and reconstruction efficiency have been further improved by Liu
et al. (2023a). Make-It-3D (Tang et al., 2023) uses BLIP to generate text descriptions for the input
image (which is applied to guide the text-to-image diffusion) and trains the model with score distilla-
tion sampling loss (Poole et al., 2022) and CLIP image loss to create geometrically and semantically
plausible shapes.
1Five seconds per shape on a single NVIDIA A100 GPU, including around 1.14 seconds image-to-triplane
feed-forward time, 1.14 seconds to query resolution of 384ˆ384ˆ384 points from the triplane-NeRF, and 1.91
seconds mesh extraction time using Marching Cubes (Lorensen & Cline, 1998).
2
In contrast to all these methods, our LRM is a purely data-driven approach that learns to reconstruct
arbitrary objects in the wild. It is trained with minimal and extensible 3D supervision (i.e., rendered
or captured 2D images of 3D objects) and does not rely on any guidance from pre-trained vision-
language contrastive or generative models.
Learning 3D Representations from Images
3D reconstruction from a single image is an ill-posed
problem that has been frequently addressed by models with generative properties. Many previous
works apply an encoder-decoder framework to model the image-to-3D data distribution (Choy et al.,
2016; Yan et al., 2016; Dai et al., 2017; Xu et al., 2019; Wu et al., 2020; M¨
uller et al., 2022; Sajjadi
et al., 2022; Goel et al., 2023), where a compact latent code is trained to carry the texture, geometry,
and pose details of the target. However, learning such an expressive representation usually requires
a capable network and abundant 3D data which is very expensive to acquire. Hence most of these
methods only focus on a few categories and produce very coarse results. GINA-3D (Shen et al.,
2023a) implements a model that applies a visual transformer encoder and cross-attention (instead
of a transformer decoder as in LRM) to translate images to triplane representations. However, the
model and training are much smaller in scale, and their work has a different focus on category-
specific generation. Recent data-driven approach MCC (Wu et al., 2023) trains a generalizable
transformer-based decoder with CO3D-v2 data (Reizenstein et al., 2021) to predict occupancy and
color from the input image and its unprojected point cloud. Although MCC can handle real and
generated images and scenes, the results are usually over-smooth and lose details.
Multimodal 3D
Motivated by the great advances in 2D multimodal learning (Tan & Bansal, 2019;
Chen et al., 2020; 2022b; Yu et al., 2022; Singh et al., 2022; Wang et al., 2022; Alayrac et al., 2022;
Girdhar et al., 2023), LRM considers 3D as a new modality and directly grounds 2D feature maps
onto 3D triplane via cross-attention. There are early attempts in this direction that minimize the dif-
ference between encoded image and 3D representations (Girdhar et al., 2016; Mandikal et al., 2018),
as well as recent research, ULIP (Xue et al., 2023) and CLIP2 (Zeng et al., 2023), which bridges
3D, language, and images via contrastive learning. LERF (Kerr et al., 2023) learns a language
field inside NeRF by rendering CLIP embeddings along training rays. In contrast, our method fo-
cuses on generic single image-to-3D reconstruction. We would like to mention the concurrent work
Cap3D (Luo et al., 2023) that produces descriptions for 3D shapes by applying BLIP (Li et al.,
2023b) to generate captions of different views, uses GPT-4 (OpenAI, 2023) to summarize them, and
then employs these language-3D pairs for training text-to-3D generative models (Nichol et al., 2022;
Poole et al., 2022; Jun & Nichol, 2023). There are also recent works in connecting 3D and large
language models, such as 3D-LLM (Hong et al., 2023) and LLM-Grounder (Yang et al., 2023).
3
METHOD
In this section, we detail the proposed LRM architecture (Fig. 1). LRM contains an image encoder
that encodes the input image to patch-wise feature tokens (Sec. 3.1), followed by an image-to-
triplane decoder that projects image features onto triplane tokens via cross-attention (Sec. 3.2). The
output triplane tokens are upsampled and reshaped into the final triplane representation, which is
used to query 3D point features. Lastly, the 3D point features are passed to a multi-layer perception
to predict RGB and density for volumetric rendering (Sec. 3.3). The training objectives and data are
described in Sec. 3.4 and Sec. 4.1.
3.1
IMAGE ENCODER
Given an RGB image as input, LRM first applies a pre-trained visual transformer (ViT) (Dosovitskiy
et al., 2020) to encode the image to patch-wise feature tokens thiun
i“1 P RdE, where i denotes the
i-th image patch, n is the total number of patches, and dE is the latent dimension of the encoder.
Specifically, we use DINO (Caron et al., 2021), a model trained with self-distillation that learns inter-
pretable attention over the structure and texture of the salient content in images. Compared to other
semantic-oriented representations such as the visual features from ImageNet-pretrained ResNet (He
et al., 2016) or CLIP (Radford et al., 2021), the detailed structural and texture information in DINO is
more important in our case since LRM can use it to reconstruct the geometry and color in 3D space.
3
Self
MLP
+
+
MLP
+
Self
+
Cross
+
Learnable
positional embeddings
Dim: (3 x 32 x 32) x 1024
Image encoder
12 Layers, Dim: 768, ViT (DINO)
Image-to-Triplane Decoder
16 Layers, Dim: 1024
Reshape
Triplane | Dim: 3 x (64 x 64) x 80
Single input image
Dim: 512 x 512 x 3
Image features
Dim: (32 x 32) x 768
Point features
Dim: 3 x 80
MLP
RGB, 𝜎
Volumetric Rendering
Rendered
novel image
Camera features, Dim: 20
Neural Radiance Field (NeRF)
10 layers, Dim 64
DeConv
Triplane tokens
Dim: (3 x 32 x 32) x 1024
Res: (3 x 32 x 32) → (3 x 64 x 64)
Dim: 1024 → 80
M
M
M
M
+
Modulation with
camera features
Residual connection
Conv
Figure 1: The overall architecture of LRM, a fully-differentiable transformer-based encoder-decoder
framework for single-image to NeRF reconstruction.
LRM applies a pre-trained vision model
(DINO) to encode the input image (Sec. 3.1), where the image features are projected to a 3D tri-
plane representation by a large transformer decoder via cross-attention (Sec. 3.2), followed by a
multi-layer perceptron to predict the point color and density for volumetric rendering (Sec. 3.3).
The entire network is trained end-to-end on around a million of 3D data (Sec. 4.1) with simple im-
age reconstruction losses (Sec. 3.4).
As a result, instead of only using the ViT pre-defined class token [CLS] that aggregates patch-wise
features, we also utilize the entire feature sequence thiun
i“1 to better preserve this information2.
3.2
IMAGE-TO-TRIPLANE DECODER
We implement a transformer decoder to project image and camera features onto learnable spatial-
positional embeddings and translate them to triplane representations. This decoder can be considered
as a prior network that is trained with large-scale data to provide necessary geometric and appearance
information to compensate for the ambiguities of single-image reconstruction.
Camera Features
We construct the camera feature c P R20 of the input image by flattening out the
4-by-4 camera extrinsic matrix E (that represents the camera-to-world transformation) and concate-
nate it with the camera focal length foc and principal point pp as c “ rE1ˆ16, focx, focy, ppx, ppys.
Moreover, we normalize the camera extrinsic E by similarity transformations so that all the input
cameras are aligned on the same axis (with the lookup direction aligned with the z-axis). Note that,
LRM does not depend on a canonical pose of the object, and the ground truth c is only applied in
training. Conditioning on normalized camera parameters greatly reduces the optimization space of
triplane features and facilitates model convergence (see details in Sec. 4.2). To embed the camera
feature, we further implement a multi-layer perceptron (MLP) to map the camera feature to a high-
dimensional camera embedding ˜
c. The intrinsics (focal and principal point) are normalized by the
image’s height and width before sending to the MLP layer.
Triplane Representation
We follow previous works (Chan et al., 2022; Gao et al., 2022) to apply
triplane as a compact and expressive feature representation of the reconstruction subject. A triplane
T contains three axis-aligned feature planes TXY , TYZ and TXZ. In our implementation, each plane
is of dimension p64ˆ64qˆdT where 64ˆ64 is the spatial resolution, and dT is the number of feature
channels. For any 3D point in the NeRF object bounding box r´1, 1s3, we can project it onto each
2For simplicity, we use thiun
i“1 in the following to denote the concatenated sequence of the encoded [CLS]
token and patch-wise features.
4
of the planes and query the corresponding point features pTxy, Tyz, Txzq via bilinear interpolation,
which is then decoded by an MLPnerf into the NeRF color and density (Sec. 3.3).
To obtain the triplane representation T , we define learnable spatial-positional embeddings f init
of dimension p3ˆ32ˆ32qˆdD which guide the image-to-3D projection and are used to query the
image features via cross-attention, where dD is the hidden dimension of the transformer decoder.
The number of tokens in f init is smaller than the number of final triplane tokens (3ˆ64ˆ64); we
will upsample the output of the transformer f out to the final T . In the forward pass, conditioning on
the camera features ˜
c and image features thiun
i“1, each layer of our image-to-triplane transformer
decoder gradually updates the initial positional embedding f init to the final triplane features via
modulation and cross-attention, respectively. The reason for applying two different conditional
operations is that the camera controls the orientation and distortion of the whole shape, whereas
the image features carry the fine-grained geometric and color information that need to be embedded
onto the triplane. Details of the two operations are explained below.
Modulation with Camera Features
Our camera modulation is inspired by DiT (Peebles & Xie,
2022) which implements an adaptive layer norm (adaLN) to modulate image latents with denoising
timesteps and class labels. Suppose tfju is a sequence of vectors in transformer, we define our
modulation function ModLNcpfjq with camera feature c as
γ, β “ MLPmodp˜
cq
(1)
ModLNcpfjq “ LNpfjq ¨ p1 ` γq ` β
(2)
where γ and β are the scale and shift (Huang & Belongie, 2017) output by MLPmod and LN is the
Layer Normalization (Ba et al., 2016). Such modulation is applied to each attention sub-layer which
will be specified next.
Transformer Layers
Each transformer layer contains a cross-attention sub-layer, a self-attention
sub-layer, and a multi-layer perceptron sub-layer (MLP), where the input tokens to each sub-layer
are modulated by the camera features. Suppose feature sequence f in is the input of an transformer
layer, we can consider f in as the triplane hidden features since they are corresponding to the final
triplane features T . As shown in the decoder part of Fig. 1, the cross-attention module firstly attends
from the triplane hidden features f in to the image features thiun
i“1, which can help linking image
information to the triplane. Note that we here do not explicitly define any spatial alignment between
the 2D images and 3D triplane hidden features, but consider 3D as an independent modality and ask
the model to learn the 2D-to-3D correspondence by itself. The updated triplane hidden features will
be passed to a self-attention sub-layer that further models the intra-modal relationships across the
spatially-structured triplane entries. Then, a multi-layer perceptron sub-layer (MLPtfm) follows as
in the original Transformer (Vaswani et al., 2017) design. Lastly, the output triplane features f out
will become the input to the next transformer layer.
Such a design is similar to the Perceiver network (Jaegle et al., 2021) while our model maintains a
high-dimensional representation across the attention layers instead of projecting the input to a latent
bottleneck. Overall, we can express this process for each j-th triplane entry in each layer as
f cross
j
“ CrossAttnpModLNcpf in
j q; thiun
i“1q ` f in
j
(3)
f self
j
“ SelfAttnpModLNcpf cross
j
q; tModLNcpf cross
j
qujq ` f cross
j
(4)
f out
j
“ MLPtfmpModLNcpf self
j
qq ` f self
j
(5)
The ModLN operators in sub-layers (i.e., CrossAttn, SelfAttn, MLPtfm) use different set of learn-
able parameters in the layer normalization and the modulation MLPmod. We do not add additional
superscript to differentiate them for clarity.
The transformer layers are processed sequentially. After all the transformer layers, we obtain the
output triplane features f out from the last layer as the output of the decoder. This final output is
upsampled by a learnable de-convolution layer and reshaped to the final triplane representation T .
3.3
TRIPLANE-NERF
We employ the triplane-NeRF formulation (Chan et al., 2022) and implement an MLPnerf to pre-
dict RGB and density σ from the point features queried from the triplane representation T . The
5
MLPnerf contains multiple linear layers with ReLU (Nair & Hinton, 2010) activation. The output
dimension of the MLPnerf is 4 where the first three dimensions are RGB colors and the last di-
mension corresponds to the density of the field. We refer to the Appendix for the details of NeRF
volumetric rendering.
3.4
TRAINING OBJECTIVES
LRM produces the 3D shape from a single input image and leverages additional side views to guide
the reconstruction during training. For each shape in the training data, we consider pV ´1q randomly
chosen side views for supervision; we apply simple image reconstruction objectives between the V
rendered views ˆ
x and the ground-truth views xGT (include the input view and side views). More
precisely, for every input image x, we minimize:
Lreconpxq “ 1
V
V
ÿ
v“1
`
LMSEpˆ
xv, xGT
v
q ` λLLPIPSpˆ
xv, xGT
v
q
˘
(6)
where LMSE is the normalized pixel-wise L2 loss, LLPIPS is the perceptual image patch similar-
ity (Zhang et al., 2018) and λ is a customized weight coefficient.
4
EXPERIMENTS
4.1
DATA
LRM relies on abundant 3D data from Objaverse (Deitke et al., 2023) and MVImgNet (Yu et al.,
2023), consisting of synthetic 3D assets and videos of objects in the real world, respectively, to learn
a generalizable cross-shape 3D prior. For each 3D asset in Objaverse, we normalize the shape to
the box r´1, 1s3 in world space and render 32 random views with the same camera pointing toward
the shape at arbitrary poses. The rendered images are of resolution 1024ˆ1024, and the camera
poses are sampled from a ball of radius r1.5, 3.0s and with height in range r´0.75, 1.60s3. For each
video, we utilize the extracted frames from the dataset. Since the target shape in those frames can
be at random positions, we crop and resize all of them using the predicted object mask4 so that the
object is at the center of the resulting frames; we adjust the camera parameters accordingly. Note
that our method does not model background, hence we render images from Objaverse with a pure
white background, and use an off-the-shelf package4 to remove the background of video frames. In
total, we pre-processed 730,648 3D assets and 220,219 videos for training.
To evaluate the performance of LRM on arbitrary images, we collected novel images from Ob-
javerse (Deitke et al., 2023), MvImgNet (Yu et al., 2023), ImageNet (Deng et al., 2009), Google
Scanned Objects (Downs et al., 2022), Amazon Berkeley Objects (Collins et al., 2022), captured
new images in the real world, and generated images with Adobe Firefly5 for reconstruction. We
visualize their results in Sec. 4.3.1 and Appendix. To numerically study the design choices of our
approach, we randomly acquired 50 unseen 3D shapes from the Objaverse and 50 unseen videos
from the MvImgNet dataset, respectively. For each shape, we pre-process 15 reference views and
pass five of them to our model one by one to reconstruct the same object, and evaluate the rendered
images using all 15 reference views (see analyses in Appendix).
4.2
IMPLEMENTATION DETAILS
Camera Normalization
We normalize the camera poses corresponding to the input images to
facilitate the image-to-triplane modeling. Specifically, for the images rendered from synthetic 3D
assets in Objaverse, regardless of the corresponding positions of the cameras, we normalize the input
camera poses to position r0, ´2, 0s with the camera vertical axis aligned with the upward z-axis in
the world frame. For the video data, since the camera can be at an arbitrary distance from the target
and the object is not at the image center, we only normalize the camera pose to r0, ´dis, 0s where
dis is the original distance between world origin and camera origin.
3Most of Objaverse assets have consistent z-axis up.
4Rembg package, a tool to remove image background: https://pypi.org/project/rembg
5Adobe Firefly, a text-to-image generation tool: https://firefly.adobe.com
6
Phone
Captured
Phone
Captured
Phone
Captured
Phone
Captured
Generated
Generated
Generated
Generated
Input Image
Rendered Novel Views
Input Image
Rendered Novel Views
ImageNet
ImageNet
Input Image
Rendered
Input Image
GT
Rendered
GT
Rendered
GT
Rendered
GT
Objaverse
Objaverse
Objaverse
Objaverse
Figure 2: Rendered novel views (RGB and depth) of shapes reconstructed by our LRM from single
images. None of the images are observed by the model during training. Generated images are
created using Adobe Firefly. The last two rows compare our results to the rendered ground truth
images of Objaverse objects (GT). Please zoom in for clearer visualization.
7
Input Image
Ours
One-2-3-45
Input Image
Ours
One-2-3-45
Figure 3: Comparison to One-2-3-45 (Liu et al., 2023a). To avoid cherry-picking, input images in the
first three rows are selected from the examples provided in One-2-3-45’s paper or demo page. None
of the images are observed by our model during training. Please zoom in for clearer visualization.
Input Image
Rendered Novel Views
Input Image
Rendered Novel Views
Input Image
Rendered Novel Views
Figure 4: Failure cases of our method. All three examples show blurry textures for occluded regions,
and distortion due to the largely inaccurate assumption of the camera parameters.
Network Architecture
We apply the ViT-B/16 model of pre-trained DINO as the image encoder,
which takes 512ˆ512 RGB images as input and produces 1025 feature tokens (1024 patch-wise
features plus one [CLS] features) of dimension 768 (dE) (Caron et al., 2021). The image-to-
triplane decoder and the MLPnerf are of 16 and 10 layers with hidden dimensions 1024 (dD) and
64, respectively. The triplane dimension is 80 (dT ). For neural rendering, LRM uniformly samples
128 points for each ray and renders 128ˆ128 resolution images for supervision. We also use the
deferred back-propagation introduced in ARF (Zhang et al., 2022) to save GPU memory.
Training
We train LRM on 128 NVIDIA (40G) A100 GPUs with batch size 1024 (1024 different
shapes per iteration) for 30 epochs, taking about 3 days to complete. Each epoch contains one
copy of the rendered image data from Objaverse and three copies of the video frame data from
MvImgNet to balance the amount of synthetic and real data. For each sample, we use 3 randomly
chosen side views (i.e., the total views V “ 4) to supervise the shape reconstruction, and we set the
coefficient λ“2.0 for LLPIPS. We apply the AdamW optimizer (Loshchilov & Hutter, 2017) and set
the learning rate to 4ˆ10´4 with a cosine schedule (Loshchilov & Hutter, 2016). We numerically
analyze the influence of data, training, and model hyper-parameters in the Appendix.
Inference
During inference, LRM takes an arbitrary image as input (squared and background re-
moved) and assumes the unknown camera parameters to be the normalized cameras that we applied
to train the Objaverse data. We query a resolution of 384ˆ384ˆ384 points from the reconstructed
triplane-NeRF and extract the mesh using Marching Cubes (Lorensen & Cline, 1998). This entire
process only takes less than 5 seconds to complete on a single NVIDIA A100 GPU.
4.3
RESULTS
We visualize the novel views of shapes reconstructed from real, generated, and rendered images from
various datasets (Fig. 2), compare our method with a concurrent work (Liu et al., 2023a) (Fig. 3), and
summarize some failure cases of our method (Sec. 4.3.2). Numerical comparisons to other methods,
and analyses of data, model architecture, and supervision can be found in the Appendix.
8
4.3.1
VISUALIZATION
Figure 2 visualizes some examples of the shapes reconstructed from single images. Overall, the
results show very high fidelity for diverse inputs, including real, generated, and rendered images of
various subjects with distinct textures. Not only is complex geometry correctly modeled (e.g. flower,
flagon, and wipe), but also the high-frequency details, such as the texture of the wood peafowl, are
preserved, both reflecting the great generalization ability of our model. From the asymmetric exam-
ples, giraffe, penguin, and bear, we can see that LRM can infer semantically reasonable occluded
portion of the shapes, which implies effective cross-shape priors have been learned.
In Figure 3, we compare LRM with One-2-3-45, a concurrent work to ours that achieves state-
of-the-art single image to 3D reconstruction by generating multi-view images with 2D diffusion
models (Liu et al., 2023a). To avoid cherry-picking, we directly test our method on the example
images provided in their paper or demo page6. We can see that our method produces much sharper
details and consistent surfaces. In the last row of the figure, we test One-2-3-45 with two examples
used in Figure 2, showing much worse reconstruction results.
4.3.2
LIMITATIONS
Despite the high-quality single-image-to-3D results we have shown, our method still has a few lim-
itations. First, our LRM tends to produce blurry textures for occluded regions, as shown in Figure
4. We conjecture that this is due to the fact that the single-image-to-3D problem is inherently prob-
abilistic, i.e., multiple plausible solutions exist for the unseen region, but our model is deterministic
and is likely producing averaged modes of the unseens. Second, during inference time, we assign a
set of fixed camera intrinsics and extrinsics (same as our Objaverse training data) to the test images.
These camera parameters may not align well with the ground truth, especially when the images are
cropped and resized, causing large changes to Field-of-View (FoV) and principal points. Figure 4
shows that incorrect assumptions of the camera parameters can lead to distorted shape reconstruc-
tion. Third, we only address images of objects without background; handling the background (Zhang
et al., 2020; Barron et al., 2022), as well as complex scenes, is beyond the scope of this work. Finally,
we assume Lambertian objects and omit the view-dependent modelling (Mildenhall et al., 2021) in
our predicted NeRF. Therefore, we cannot faithfully reconstruct the view-dependent appearance of
some real-world materials, e.g., shiny metals, glossy ceramics, etc.
5
CONCLUSION
In this paper, we propose LRM, the first large transformer-based framework to learn an expressive
3D prior from a million 3D data to reconstruct objects from single images. LRM is very efficient
in training and inference; it is a fully-differentiable network that can be trained end-to-end with
simple image reconstruction losses and only takes five seconds to render a high-fidelity 3D shape,
thus enabling a wide range of real-world applications. In the era of large-scale learning, we hope our
idea can inspire future research to explore data-driven 3D large reconstruction models that generalize
well to arbitrary in-the-wild images.
Future Directions
In addition to addressing the limitations mentioned in Sec. 4.3.2, we suggest
two future directions of our research; (1) Scaling up the model and training data: with the simplest
transformer-based design and minimal regularization, LRM can be easily scaled to a larger and more
capable network, including but not limited to applying a larger image encoder, adding more attention
layers to the image-to-triplane decoder, and increasing the resolution of triplane representations. On
the other hand, LRM only requires multi-view images for supervision, hence a wide range of 3D,
video, and image datasets can be exploited in training. We expect both approaches to be promising
in improving the model’s generalization ability and the quality of reconstruction. (2) Extension to
multimodal 3D generative models: LRM model builds a pathway for generating novel 3D shapes
from language by leveraging a text-to-image generation model to first create 2D images. But more
interestingly, we suggest the learned expressive triplane representations could be applied to directly
bridge language descriptions and 3D to enable efficient text-to-3D generation and editing (e.g., via
latent diffusion (Rombach et al., 2022)). We will explore these ideas in our future research.
6One-2-3-45 demo page: https://huggingface.co/spaces/One-2-3-45/One-2-3-45.
9
ETHICS STATEMENT
LRM proposed in this paper is a deterministic model in which, given the same image as input, the
model will infer the identical 3D shape. Unlike generative models that can be used to easily synthe-
size various undesirable contents (e.g., from language inputs), LRM requests the specific 2D content
to exist in the first place. LRM is trained on Objaverse (Deitke et al., 2023) and MvImgNet (Yu et al.,
2023) data, which mostly contain ethical content. However, given an unethical or misleading image,
LRM could produce unethical 3D objects or 3D disinformation that may be more convincing than
the 2D input images (although the reconstructed objects are less realistic than real-world objects).
Image-to-3D reconstruction models like LRM hold the potential to automate tasks currently per-
formed by 3D designers. However, it’s worth noting that these tools also have the capacity to foster
growth and enhance accessibility within the creative industry.
REPRODUCIBILITY STATEMENT
Our LRM is built by integrating the publicly available codebases of threestudio7 (Guo et al., 2023),
x-transformers8, and DINO9 (Caron et al., 2021), and the model is trained using publicly available
data from Objaverse (Deitke et al., 2023) and MvImgNet (Yu et al., 2023). We include very compre-
hensive data pre-processing, network architecture, and training details in this paper, which greatly
facilitate reproducing our LRM.
ACKNOWLEDGMENT
We want to thank Nathan Carr, Scott Cohen, Hailin Jin, Aseem Agarwala, Tong Sun for their sup-
port, and thank Duygu Ceylan, Zexiang Xu, Paul Guerrero, Chun-Hao Huang, Niloy Mitra, Radomir
Mech, Vova Kim, Thibault Groueix for constructive feedback on this project. Hao wants to thank
Xin for the inspiration as he ran on this road. Yicong wants to thank Prof. Stephen Gould and Ms.
Ziwei Wang for their great advice.
REFERENCES
Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel
Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language
model for few-shot learning. Advances in Neural Information Processing Systems, 35:23716–
23736, 2022.
Titas Anciukeviˇ
cius, Zexiang Xu, Matthew Fisher, Paul Henderson, Hakan Bilen, Niloy J Mitra, and
Paul Guerrero. Renderdiffusion: Image diffusion for 3d reconstruction, inpainting and generation.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
12608–12618, 2023.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint
arXiv:1607.06450, 2016.
Jonathan T. Barron, Ben Mildenhall, Dor Verbin, Pratul P. Srinivasan, and Peter Hedman. Mip-nerf
360: Unbounded anti-aliased neural radiance fields. CVPR, 2022.
Federica Bogo, Angjoo Kanazawa, Christoph Lassner, Peter Gehler, Javier Romero, and Michael J
Black. Keep it smpl: Automatic estimation of 3d human pose and shape from a single image. In
Computer Vision–ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, Octo-
ber 11-14, 2016, Proceedings, Part V 14, pp. 561–578. Springer, 2016.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are
few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
7threestudio’s GitHub page: https://github.com/threestudio-project/threestudio.
8x-transformers’s GitHub page: https://github.com/lucidrains/x-transformers.
9DINO’s GitHub page: https://github.com/facebookresearch/dino.
10
Mathilde Caron, Hugo Touvron, Ishan Misra, Herv´
e J´
egou, Julien Mairal, Piotr Bojanowski, and
Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of
the IEEE/CVF international conference on computer vision, pp. 9650–9660, 2021.
Eric R Chan, Connor Z Lin, Matthew A Chan, Koki Nagano, Boxiao Pan, Shalini De Mello, Orazio
Gallo, Leonidas J Guibas, Jonathan Tremblay, Sameh Khamis, et al. Efficient geometry-aware
3d generative adversarial networks. In Proceedings of the IEEE/CVF Conference on Computer
Vision and Pattern Recognition, pp. 16123–16133, 2022.
Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li,
Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al. Shapenet: An information-rich 3d
model repository. arXiv preprint arXiv:1512.03012, 2015.
Anpei Chen, Zexiang Xu, Andreas Geiger, Jingyi Yu, and Hao Su. Tensorf: Tensorial radiance
fields. In European Conference on Computer Vision (ECCV), 2022a.
Jun Chen, Han Guo, Kai Yi, Boyang Li, and Mohamed Elhoseiny. Visualgpt: Data-efficient adap-
tation of pretrained language models for image captioning. In Proceedings of the IEEE/CVF
Conference on Computer Vision and Pattern Recognition, pp. 18030–18040, 2022b.
Yen-Chun Chen, Linjie Li, Licheng Yu, Ahmed El Kholy, Faisal Ahmed, Zhe Gan, Yu Cheng, and
Jingjing Liu. Uniter: Universal image-text representation learning. In European conference on
computer vision, pp. 104–120. Springer, 2020.
Zhiqin Chen and Hao Zhang. Learning implicit fields for generative shape modeling. In Proceedings
of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 5939–5948, 2019.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam
Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm:
Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
Christopher B Choy, Danfei Xu, JunYoung Gwak, Kevin Chen, and Silvio Savarese. 3d-r2n2: A
unified approach for single and multi-view 3d object reconstruction. In Computer Vision–ECCV
2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceed-
ings, Part VIII 14, pp. 628–644. Springer, 2016.
Jasmine Collins, Shubham Goel, Kenan Deng, Achleshwar Luthra, Leon Xu, Erhan Gundogdu,
Xi Zhang, Tomas F Yago Vicente, Thomas Dideriksen, Himanshu Arora, et al. Abo: Dataset and
benchmarks for real-world 3d object understanding. In Proceedings of the IEEE/CVF Conference
on Computer Vision and Pattern Recognition, pp. 21126–21136, 2022.
Angela Dai, Charles Ruizhongtai Qi, and Matthias Nießner. Shape completion using 3d-encoder-
predictor cnns and shape synthesis. In Proceedings of the IEEE conference on computer vision
and pattern recognition, pp. 5868–5877, 2017.
Matt Deitke, Dustin Schwenk, Jordi Salvador, Luca Weihs, Oscar Michel, Eli VanderBilt, Ludwig
Schmidt, Kiana Ehsani, Aniruddha Kembhavi, and Ali Farhadi. Objaverse: A universe of anno-
tated 3d objects. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pp. 13142–13153, 2023.
Congyue Deng, Chiyu Jiang, Charles R Qi, Xinchen Yan, Yin Zhou, Leonidas Guibas, Dragomir
Anguelov, et al. Nerdi: Single-view nerf synthesis with language-guided diffusion as general
image priors.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pp. 20637–20647, 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi-
erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition,
pp. 248–255. Ieee, 2009.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
11
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas
Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An
image is worth 16x16 words: Transformers for image recognition at scale.
arXiv preprint
arXiv:2010.11929, 2020.
Laura Downs, Anthony Francis, Nate Koenig, Brandon Kinman, Ryan Hickman, Krista Reymann,
Thomas B McHugh, and Vincent Vanhoucke. Google scanned objects: A high-quality dataset
of 3d scanned household items. In 2022 International Conference on Robotics and Automation
(ICRA), pp. 2553–2560. IEEE, 2022.
Haoqiang Fan, Hao Su, and Leonidas J Guibas. A point set generation network for 3d object recon-
struction from a single image. In Proceedings of the IEEE conference on computer vision and
pattern recognition, pp. 605–613, 2017.
Jun Gao, Tianchang Shen, Zian Wang, Wenzheng Chen, Kangxue Yin, Daiqing Li, Or Litany, Zan
Gojcic, and Sanja Fidler. Get3d: A generative model of high quality 3d textured shapes learned
from images. Advances In Neural Information Processing Systems, 35:31841–31854, 2022.
Rohit Girdhar, David F Fouhey, Mikel Rodriguez, and Abhinav Gupta. Learning a predictable and
generative vector representation for objects. In Computer Vision–ECCV 2016: 14th European
Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part VI 14, pp.
484–499. Springer, 2016.
Rohit Girdhar, Alaaeldin El-Nouby, Zhuang Liu, Mannat Singh, Kalyan Vasudev Alwala, Armand
Joulin, and Ishan Misra. Imagebind: One embedding space to bind them all. In Proceedings of
the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15180–15190, 2023.
Georgia Gkioxari, Jitendra Malik, and Justin Johnson. Mesh r-cnn. In Proceedings of the IEEE/CVF
international conference on computer vision, pp. 9785–9795, 2019.
Shubham Goel, Angjoo Kanazawa, and Jitendra Malik. Shape and viewpoint without keypoints. In
Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020,
Proceedings, Part XV 16, pp. 88–104. Springer, 2020.
Shubham Goel, Georgios Pavlakos, Jathushan Rajasegaran, Angjoo Kanazawa, and Jitendra Malik.
Humans in 4d: Reconstructing and tracking humans with transformers. In Proceedings of the
IEEE/CVF International Conference on Computer Vision, pp. 14783–14794, 2023.
Yuan-Chen Guo, Ying-Tian Liu, Ruizhi Shao, Christian Laforte, Vikram Voleti, Guan Luo, Chia-
Hao Chen, Zi-Xin Zou, Chen Wang, Yan-Pei Cao, and Song-Hai Zhang. threestudio: A unified
framework for 3d content generation. https://github.com/threestudio-project/
threestudio, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog-
nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp.
770–778, 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2023.
Yining Hong, Haoyu Zhen, Peihao Chen, Shuhong Zheng, Yilun Du, Zhenfang Chen, and Chuang
Gan. 3d-llm: Injecting the 3d world into large language models. arXiv, 2023.
Xun Huang and Serge Belongie. Arbitrary style transfer in real-time with adaptive instance normal-
ization. In Proceedings of the IEEE international conference on computer vision, pp. 1501–1510,
2017.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira.
Perceiver: General perception with iterative attention. In International conference on machine
learning, pp. 4651–4664. PMLR, 2021.
Wonbong Jang and Lourdes Agapito.
Codenerf: Disentangled neural radiance fields for object
categories. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp.
12949–12958, 2021.
12
Heewoo Jun and Alex Nichol. Shap-e: Generating conditional 3d implicit functions. arXiv preprint
arXiv:2305.02463, 2023.
Angjoo Kanazawa, Shubham Tulsiani, Alexei A Efros, and Jitendra Malik. Learning category-
specific mesh reconstruction from image collections. In Proceedings of the European Conference
on Computer Vision (ECCV), pp. 371–386, 2018.
Justin Kerr, Chung Min Kim, Ken Goldberg, Angjoo Kanazawa, and Matthew Tancik. Lerf: Lan-
guage embedded radiance fields. arXiv preprint arXiv:2303.09553, 2023.
Nilesh Kulkarni, Abhinav Gupta, David F Fouhey, and Shubham Tulsiani.
Articulation-aware
canonical surface mapping. In Proceedings of the IEEE/CVF Conference on Computer Vision
and Pattern Recognition, pp. 452–461, 2020.
Jiahao Li, Hao Tan, Kai Zhang, Zexiang Xu, Fujun Luan, Yinghao Xu, Yicong Hong, Kalyan
Sunkavalli, Greg Shakhnarovich, and Sai Bi. Instant3d: Fast text-to-3d with sparse-view gen-
eration and large reconstruction model. arXiv preprint arXiv:2311.06214, 2023a.
Junnan Li, Dongxu Li, Caiming Xiong, and Steven Hoi. Blip: Bootstrapping language-image pre-
training for unified vision-language understanding and generation. In International Conference
on Machine Learning, pp. 12888–12900. PMLR, 2022.
Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi.
Blip-2:
Bootstrapping language-
image pre-training with frozen image encoders and large language models.
arXiv preprint
arXiv:2301.12597, 2023b.
Xueting Li, Sifei Liu, Kihwan Kim, Shalini De Mello, Varun Jampani, Ming-Hsuan Yang, and
Jan Kautz. Self-supervised single-view 3d reconstruction via semantic consistency. In Computer
Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings,
Part XIV 16, pp. 677–693. Springer, 2020.
Minghua Liu, Chao Xu, Haian Jin, Linghao Chen, Zexiang Xu, Hao Su, et al.
One-2-3-45:
Any single image to 3d mesh in 45 seconds without per-shape optimization.
arXiv preprint
arXiv:2306.16928, 2023a.
Ruoshi Liu, Rundi Wu, Basile Van Hoorick, Pavel Tokmakov, Sergey Zakharov, and Carl Vondrick.
Zero-1-to-3: Zero-shot one image to 3d object. arXiv preprint arXiv:2303.11328, 2023b.
William E Lorensen and Harvey E Cline. Marching cubes: A high resolution 3d surface construction
algorithm. In Seminal graphics: pioneering efforts that shaped the field, pp. 347–353. 1998.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv
preprint arXiv:1608.03983, 2016.
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization.
arXiv preprint
arXiv:1711.05101, 2017.
Tiange Luo, Chris Rockwell, Honglak Lee, and Justin Johnson. Scalable 3d captioning with pre-
trained models. arXiv preprint arXiv:2306.07279, 2023.
Priyanka Mandikal, KL Navaneet, Mayank Agarwal, and R Venkatesh Babu. 3d-lmnet: Latent
embedding matching for accurate and diverse 3d point cloud reconstruction from a single image.
arXiv preprint arXiv:1807.07796, 2018.
Luke Melas-Kyriazi, Iro Laina, Christian Rupprecht, and Andrea Vedaldi.
Realfusion: 360deg
reconstruction of any object from a single image. In Proceedings of the IEEE/CVF Conference
on Computer Vision and Pattern Recognition, pp. 8446–8455, 2023.
Lars Mescheder, Michael Oechsle, Michael Niemeyer, Sebastian Nowozin, and Andreas Geiger. Oc-
cupancy networks: Learning 3d reconstruction in function space. In Proceedings of the IEEE/CVF
conference on computer vision and pattern recognition, pp. 4460–4470, 2019.
Gal Metzer, Elad Richardson, Or Patashnik, Raja Giryes, and Daniel Cohen-Or. Latent-nerf for
shape-guided generation of 3d shapes and textures. In Proceedings of the IEEE/CVF Conference
on Computer Vision and Pattern Recognition, pp. 12663–12673, 2023.
13
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and
Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications
of the ACM, 65(1):99–106, 2021.
Paritosh Mittal, Yen-Chi Cheng, Maneesh Singh, and Shubham Tulsiani. Autosdf: Shape priors for
3d completion, reconstruction and generation. In Proceedings of the IEEE/CVF Conference on
Computer Vision and Pattern Recognition, pp. 306–315, 2022.
Norman M¨
uller, Andrea Simonelli, Lorenzo Porzi, Samuel Rota Bulo, Matthias Nießner, and Peter
Kontschieder. Autorf: Learning 3d object radiance fields from single view observations. In Pro-
ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3971–
3980, 2022.
Thomas M¨
uller, Alex Evans, Christoph Schied, and Alexander Keller. Instant neural graphics prim-
itives with a multiresolution hash encoding.
ACM Trans. Graph., 41(4):102:1–102:15, July
2022.
doi: 10.1145/3528223.3530127.
URL https://doi.org/10.1145/3528223.
3530127.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In
Proceedings of the 27th international conference on machine learning (ICML-10), pp. 807–814,
2010.
Alex Nichol, Heewoo Jun, Prafulla Dhariwal, Pamela Mishkin, and Mark Chen. Point-e: A system
for generating 3d point clouds from complex prompts. arXiv preprint arXiv:2212.08751, 2022.
Michael Niemeyer, Lars Mescheder, Michael Oechsle, and Andreas Geiger. Differentiable volumet-
ric rendering: Learning implicit 3d representations without 3d supervision. In Proceedings of the
IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3504–3515, 2020.
David Novotny, Nikhila Ravi, Benjamin Graham, Natalia Neverova, and Andrea Vedaldi. C3dpo:
Canonical 3d pose networks for non-rigid structure from motion. In Proceedings of the IEEE/CVF
International Conference on Computer Vision, pp. 7688–7697, 2019.
OpenAI. Gpt-4 technical report, 2023.
Jeong Joon Park, Peter Florence, Julian Straub, Richard Newcombe, and Steven Lovegrove.
Deepsdf: Learning continuous signed distance functions for shape representation. In Proceedings
of the IEEE/CVF conference on computer vision and pattern recognition, pp. 165–174, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor
Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-
performance deep learning library. Advances in neural information processing systems, 32, 2019.
William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint
arXiv:2212.09748, 2022.
Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d
diffusion. arXiv preprint arXiv:2209.14988, 2022.
Guocheng Qian, Jinjie Mai, Abdullah Hamdi, Jian Ren, Aliaksandr Siarohin, Bing Li, Hsin-
Ying Lee, Ivan Skorokhodov, Peter Wonka, Sergey Tulyakov, et al.
Magic123: One image
to high-quality 3d object generation using both 2d and 3d diffusion priors.
arXiv preprint
arXiv:2306.17843, 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language
models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal,
Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual
models from natural language supervision. In International conference on machine learning, pp.
8748–8763. PMLR, 2021.
14
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen,
and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine
Learning, pp. 8821–8831. PMLR, 2021.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-
conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
Jeremy Reizenstein, Roman Shapovalov, Philipp Henzler, Luca Sbordone, Patrick Labatut, and
David Novotny. Common objects in 3d: Large-scale learning and evaluation of real-life 3d cat-
egory reconstruction. In Proceedings of the IEEE/CVF International Conference on Computer
Vision, pp. 10901–10911, 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bj¨
orn Ommer. High-
resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer-
ence on computer vision and pattern recognition, pp. 10684–10695, 2022.
Joseph Roth, Yiying Tong, and Xiaoming Liu. Adaptive 3d face reconstruction from unconstrained
photo collections. In Proceedings of the IEEE conference on computer vision and pattern recog-
nition, pp. 4197–4206, 2016.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar
Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic
text-to-image diffusion models with deep language understanding. Advances in Neural Informa-
tion Processing Systems, 35:36479–36494, 2022.
Mehdi SM Sajjadi, Henning Meyer, Etienne Pot, Urs Bergmann, Klaus Greff, Noha Radwan, Suhani
Vora, Mario Luˇ
ci´
c, Daniel Duckworth, Alexey Dosovitskiy, et al. Scene representation trans-
former: Geometry-free novel view synthesis through set-latent scene representations. In Proceed-
ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6229–6238,
2022.
Bokui Shen, Xinchen Yan, Charles R Qi, Mahyar Najibi, Boyang Deng, Leonidas Guibas, Yin
Zhou, and Dragomir Anguelov. Gina-3d: Learning to generate implicit neural assets in the wild.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
4913–4926, 2023a.
Qiuhong Shen, Xingyi Yang, and Xinchao Wang.
Anything-3d: Towards single-view anything
reconstruction in the wild. arXiv preprint arXiv:2304.10261, 2023b.
Amanpreet Singh, Ronghang Hu, Vedanuj Goswami, Guillaume Couairon, Wojciech Galuba, Mar-
cus Rohrbach, and Douwe Kiela. Flava: A foundational language and vision alignment model.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
15638–15650, 2022.
Cheng Sun, Min Sun, and Hwann-Tzong Chen. Direct voxel grid optimization: Super-fast conver-
gence for radiance fields reconstruction. In CVPR, 2022.
Hao Tan and Mohit Bansal. Lxmert: Learning cross-modality encoder representations from trans-
formers. arXiv preprint arXiv:1908.07490, 2019.
Junshu Tang, Tengfei Wang, Bo Zhang, Ting Zhang, Ran Yi, Lizhuang Ma, and Dong Chen.
Make-it-3d: High-fidelity 3d creation from a single image with diffusion prior. arXiv preprint
arXiv:2303.14184, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timoth´
ee
Lacroix, Baptiste Rozi`
ere, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and
efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
Shubham Tulsiani, Tinghui Zhou, Alexei A Efros, and Jitendra Malik.
Multi-view supervision
for single-view reconstruction via differentiable ray consistency. In Proceedings of the IEEE
conference on computer vision and pattern recognition, pp. 2626–2634, 2017.
15
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa-
tion processing systems, 30, 2017.
Nanyang Wang, Yinda Zhang, Zhuwen Li, Yanwei Fu, Wei Liu, and Yu-Gang Jiang. Pixel2mesh:
Generating 3d mesh models from single rgb images. In Proceedings of the European conference
on computer vision (ECCV), pp. 52–67, 2018.
Yi Wang, Kunchang Li, Yizhuo Li, Yinan He, Bingkun Huang, Zhiyu Zhao, Hongjie Zhang, Jilan
Xu, Yi Liu, Zun Wang, et al. Internvideo: General video foundation models via generative and
discriminative learning. arXiv preprint arXiv:2212.03191, 2022.
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment:
from error visibility to structural similarity. IEEE transactions on image processing, 13(4):600–
612, 2004.
Chao-Yuan Wu, Justin Johnson, Jitendra Malik, Christoph Feichtenhofer, and Georgia Gkioxari.
Multiview compressive coding for 3d reconstruction. In Proceedings of the IEEE/CVF Confer-
ence on Computer Vision and Pattern Recognition, pp. 9065–9075, 2023.
Rundi Wu, Yixin Zhuang, Kai Xu, Hao Zhang, and Baoquan Chen. Pq-net: A generative part
seq2seq network for 3d shapes. In Proceedings of the IEEE/CVF Conference on Computer Vision
and Pattern Recognition, pp. 829–838, 2020.
Dejia Xu, Yifan Jiang, Peihao Wang, Zhiwen Fan, Yi Wang, and Zhangyang Wang. Neurallift-360:
Lifting an in-the-wild 2d photo to a 3d object with 360deg views. In Proceedings of the IEEE/CVF
Conference on Computer Vision and Pattern Recognition, pp. 4479–4489, 2023.
Qiangeng Xu, Weiyue Wang, Duygu Ceylan, Radomir Mech, and Ulrich Neumann. Disn: Deep
implicit surface network for high-quality single-view 3d reconstruction. Advances in neural in-
formation processing systems, 32, 2019.
Le Xue, Mingfei Gao, Chen Xing, Roberto Mart´
ın-Mart´
ın, Jiajun Wu, Caiming Xiong, Ran Xu,
Juan Carlos Niebles, and Silvio Savarese. Ulip: Learning a unified representation of language,
images, and point clouds for 3d understanding. In Proceedings of the IEEE/CVF Conference on
Computer Vision and Pattern Recognition, pp. 1179–1189, 2023.
Xinchen Yan, Jimei Yang, Ersin Yumer, Yijie Guo, and Honglak Lee.
Perspective transformer
nets: Learning single-view 3d object reconstruction without 3d supervision. Advances in neural
information processing systems, 29, 2016.
Jianing Yang, Xuweiyi Chen, Shengyi Qian, Nikhil Madaan, Madhavan Iyengar, David F. Fouhey,
and Joyce Chai. Llm-grounder: Open-vocabulary 3d visual grounding with large language model
as an agent, 2023.
Alex Yu, Vickie Ye, Matthew Tancik, and Angjoo Kanazawa. pixelnerf: Neural radiance fields from
one or few images. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pp. 4578–4587, 2021.
Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui
Wu.
Coca:
Contrastive captioners are image-text foundation models.
arXiv preprint
arXiv:2205.01917, 2022.
Xianggang Yu, Mutian Xu, Yidan Zhang, Haolin Liu, Chongjie Ye, Yushuang Wu, Zizheng Yan,
Chenming Zhu, Zhangyang Xiong, Tianyou Liang, et al. Mvimgnet: A large-scale dataset of
multi-view images. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pp. 9150–9161, 2023.
Yihan Zeng, Chenhan Jiang, Jiageng Mao, Jianhua Han, Chaoqiang Ye, Qingqiu Huang, Dit-Yan
Yeung, Zhen Yang, Xiaodan Liang, and Hang Xu.
Clip2: Contrastive language-image-point
pretraining from real-world point cloud data. In Proceedings of the IEEE/CVF Conference on
Computer Vision and Pattern Recognition, pp. 15244–15253, 2023.
16
Kai Zhang, Gernot Riegler, Noah Snavely, and Vladlen Koltun. Nerf++: Analyzing and improving
neural radiance fields. arXiv preprint arXiv:2010.07492, 2020.
Kai Zhang, Nick Kolkin, Sai Bi, Fujun Luan, Zexiang Xu, Eli Shechtman, and Noah Snavely. Arf:
Artistic radiance fields. In European Conference on Computer Vision, pp. 717–733. Springer,
2022.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable
effectiveness of deep features as a perceptual metric. In CVPR, 2018.
17
APPENDICES
A
BACKGROUND OF MODEL COMPONENTS
A.1
NERF
We adopt NeRF (Mildenhall et al., 2021), specifically the compact triplane NeRF variant (Chan
et al., 2022), as our 3D representation to predict in LRM. NeRF, when coupled with differentiable
volume rendering, can be optimized with just image reconstruction losses.
At the core of NeRF (Mildenhall et al., 2021) and its variants (Chan et al., 2022; Chen et al., 2022a;
M¨
uller et al., 2022; Sun et al., 2022) is a spatially-varying color (modeling appearance) and density
(modeling geometry) field function. 10 Given a 3D point p, the color and density field pu, σq can be
written as:
pu, σq “ MLPnerf pfθppqq,
(7)
where the spatial encoding fθ is used to facilitate the MLPnerf to learn high-frequency signals.
Different NeRF variants (Chan et al., 2022; Chen et al., 2022a; M¨
uller et al., 2022; Sun et al., 2022)
typically differ from each other in terms of the choice of the spatial encoding and the size of the
MLP. In this work, we use the triplane spatial encoding function proposed by EG3D (Chan et al.,
2022), because of its low tokenization complexity (OpN 2q as opposed to a voxel grid’s OpN 3q
complexity, where N is spatial resolution).
Images are rendered from NeRF using volume rendering that’s trivially differentiable. In detail, for
each pixel to render, we cast a ray r through a NeRF, and use finite point samples pi along the ray
to compute the volume rendering integral to get the rendered color uprq:
uprq “
ÿ
i
Tip1 ´ expp´σiδiqqui,
(8)
Ti “ expp´
i´1
ÿ
j“1
σjδjq,
(9)
where pui, σiq “ MLPϕpfθppiqq and δi is the distance between point pi and pi`1.
A.2
TRANSFORMER LAYERS
In this subsection, we provide the details of the layers used in the transformer decoder (Vaswani
et al., 2017) as a background. For the Vision Transformer encoder, please refer to the original DINO
paper (Caron et al., 2021) for implementation details.
Attention operator
Attention operator is an expressive neural operator which converts an input
feature x with condition to a sequence of other features tyiu. It first computes the attention score αi
by using the dot product between the input x and each condition feature yi. An additional softmax
is added after the dot products to normalize the weights to a summation of 1. This attention score
measures the relationship between input and conditions. Then the output is the weighted summation
of the conditions tyiu with respect to the attention score αi.
αi “ softmaxitxJyiu
(10)
Attnpx; tyiuiq “
ÿ
i
αiyi
(11)
For some specific cases (e.g., in the transformer attention layer below), the attention operator wants
to differentiate the vectors used in calculating the attention score and the vectors for final outputs.
Thus it will introduce another set of ‘value’ vectors tziui, and treat the tyiui as corresponding ‘key’
vectors. Taking this into consideration, the formula would become
αi “ softmaxitxJyiu
(12)
Attnpx; tyiui, tziuiq “
ÿ
i
αizi
(13)
10To simplify the discussion, we ignore the view-dependent modeling in NeRF (Mildenhall et al., 2021).
18
(Q, K, V)
…
Cross-Attention
Image features
𝒉𝑖𝑖=1
𝑛
Triplane hidden
features
(K, V)
(K, V)
(K, V)
(K, V)
(Q)
(Q)
(Q)
(Q)
…
…
(Q, K, V)
(Q, K, V)
(Q, K, V)
…
Self-attended
features 𝒇𝑠𝑒𝑙𝑓
Self-Attention
Figure 5: Visual illustration of the cross-attention and self-attention in LRM’s image-to-triplane
decoder.
Multi-head Attention
The attention operator described above only attends to the condition fea-
tures once to get the attention vector. However, the actual attention might contain multiple modes.
Thus, the multi-head attention (Vaswani et al., 2017) is proposed. The multi-head attention is im-
plemented by first splitting the input features into smaller queries.
rx1, . . . , xnhs “ x
(14)
where nh is the number of heads. Meanwhile, yi and zi are split into tyk
i uk and tzk
i uk in a sim-
ilar way. After that, the output of each head is computed independently and the final output is a
concatenation of heads’ outputs.
outk “ Attnpxk; tyk
i ui, tzk
i uiq
(15)
MultiHeadAttnpx; tyiui, tziuiq “ rout1, . . . , outnhs
(16)
Attention Layers in Transformer
The detailed attention layers in transformer utilize the above
multi-head attention with more linear layers. Here are the formulas for the self-attention layer (see
the right yellow ‘Self-Attention’ block in Fig. 5). The layer first projects the input feature sequence
f “ tfjuj to query q, key k, and value v vectors with linear layers. Then the multi-head attention is
applied. There is one more linear layer over the output. We also follow the recent papers (Chowdhery
et al., 2022; Touvron et al., 2023) to remove the bias terms in the attention layers.
qj “ Wqfj
(17)
ki “ Wkfi
(18)
vi “ Wvfi
(19)
oj “ MultiHeadAttnpqj; tkiui, tviuiq
(20)
SelfAttnpfj; tfjujq “ Woutoj
(21)
(22)
The cross-attention layer is defined similarly (see the left blue ‘Cross-Attention’ block in Fig. 5).
The only difference to the self-attention layer is that the Wk and Wv is applied to the condition
vectors (e.g., the image features h in our example).
MLP layers in Transformer
The Transformer model architecture applies the MLP layer (multi-
layer perceptron) to do channel mixing (i.e., mix the information from different feature dimensions).
We follow the original transformer paper (Vaswani et al., 2017) for the implementation. The MLP
layer contains two linear layers with a GELU (Hendrycks & Gimpel, 2023) activation in between.
The intermediate hidden dimension is 4 times of the model dimension.
19
Layer Normalization
We take the default LayerNorm (LN) implementation in PyTorch (Paszke
et al., 2019). Besides the LN layers in ModLN as in Sec. 3.2, we follow the Pre-LN architecture
to also apply LN to the final output of transformers, e.g., the output of ViT and also the output of
transformer decoder.
Positional Encoding
The positional embedding in ViT (Dosovitskiy et al., 2020) is bilinearly up-
sampled from its original resolution (14ˆ14 for input 224ˆ224) to match our higher input resolution
(32ˆ32 for input 512ˆ512).
B
TRAINING SETUP
We specify the training setup of our LRM. Apart from the information that we provided in Sec. 4.2,
we apply a cosine schedule (Loshchilov & Hutter, 2016) with 3000 warm-up iterations. We set the
second beta parameter (β2) of the AdamW optimizer (Loshchilov & Hutter, 2017) to be 0.95. We
apply a gradient clipping of 1.0 and a weight decay of 0.05. The weight decay are only applied on
the weights that are not bias and not in the layer normalization layer. We use BF16 precision in in
the mixed precision training. To save computational cost in training, we resize the reference novel
views from 512ˆ512 to a randomly chosen resolution between 128ˆ128 and 384ˆ384 and only ask
the model to reconstruct a randomly selected 128ˆ128 region. With this design, we can possibly
increase the effective resolution of the model.
C
COMPARISON WITH SOTA
We provide a quantitative comparison to the stat-of-the-art methods Point-E (Nichol et al., 2022),
Shap-E (Jun & Nichol, 2023), and One-2-3-45 (Liu et al., 2023a). Point-E trains an image-to-3D
point cloud diffusion model, Shap-E encodes point clouds to latent representations and trains a
diffusion model on the latents to generate parameters of a 3D implicit function, and One-2-3-45
reconstructs multi-view images generated with a 2D diffusion model. We randomly selected 100
objects from the Google Scanned Objects (GSO) dataset (Downs et al., 2022) and measured the
novel view synthetic quality of 20 reference views (FID, CLIP-Similarity (Radford et al., 2021),
PSNR, LPIPS (Zhang et al., 2018)) and the geometric quality (Chamfer Distance), as shown in the
Table below. We can see that our LRM consistently outperforms previous approaches in all metrics.
Table 1: Comparison between LRM and state-of-the-art 3D generative models on Google Scanned
Objects dataset (100 randomly selected objects and 20 reference views).
Models
GSO Evaluation
FIDÓ
CLIP-SimilarityÒ
PSNRÒ
LPIPSÓ
Chamfer DistanceÓ
Point-E
123.70
0.741
15.60
0.308
0.099
Shap-E
97.05
0.805
14.36
0.289
0.085
One-2-3-45
139.24
0.713
12.42
0.448
0.123
LRM (ours)
31.44
0.902
19.60
0.163
0.053
We would like to discuss further the difference between LRM and the large-scale approaches Point-
E and Shap-E. The models of Point-E and Shap-E contain hundreds of millions of learnable pa-
rameters and are trained with several million 3D assets (unknown data source and unknown com-
putational cost from their papers). In terms of the network and dataset sizes, our LRM has 500
million learnable parameters, and it is trained on 1 million 3D data (publicly accessible), which
does not show an advantage. In terms of the network architecture, Point-E, Shap-E, and LRM all
use transformer-based models and apply cross-attention for inter-modality modeling (i.e., image-to-
point cloud, point cloud+image-to-3D latents, and image-to-triplane, respectively). We hypothesize
it is the choice of very compact and expressive triplane representation together with an end-to-end
trainable framework that enables the effective scaling of LRM and its adequate learning on large
datasets (Objaverse and MvImgNet). Compared to the unstructured point cloud representation ap-
plied in Point-E and Shap-E, LRM applies the structured triplane representation that is aligned with
the world frame, which naturally facilitates 2D-to-3D projection. It is also worth mentioning that
Point-E uses 4K points (as tokens) and Shap-E uses 16K points (as tokens), but our LRM only uses
20
3ˆ32ˆ32“3072 triplane tokens, which largely reduce the modeling complexity. Additionally, com-
pared to the two-stage approach in Shape-E, which attempts to generate latents that can produce the
parameters of implicit 3D functions through a diffusion model, our LRM directly maps 2D images
to triplanes, which should be much more stable and efficient to learn. Overall, we suggest that LRM
is a more data-friendly and efficient model than Point-E and Shap-E.
D
ANALYSES
We evaluate the effect of data, model hyper-parameters, and training methods on the performance
of LRM, measuring by PSNR, CLIP-Similarity (Radford et al., 2021), SSIM (Wang et al., 2004)
and LPIPS (Zhang et al., 2018) of the rendered novel views. Note that due to the large training cost
of our final model, the following analytic experiments use a much smaller version of LRM model
as the baseline (indicated by orange shaded rows in the tables). Specifically, we scale down the
image-to-triplane decoder to 12 cross-attention layers, change the input image resolution to 256,
triplane latent dimension to 32, rendering resolution in training to 64, and use 96 samples per ray for
rendering 64ˆ64 images for supervision. We only train each model on 32 NVIDIA A100 GPUs for
15 epochs, and the resulting difference can be seen in Table 2. We are aware that some observations
might change if we scale up the model, but most of the conclusions should be general and consistent.
Table 2: Comparison between the final model and the baseline for analysis.
Models
Unseen Evaluation
PSNRÒ
CLIP-SimilarityÒ
SSIMÒ
LPIPSÓ
Final
20.1
91.0
79.7
16.0
Baseline
19.0
87.8
77.4
19.1
D.1
SYNTHETIC VS. REAL DATA
Table 3 compares the influence of using synthetic 3D data from the Objaverse (Deitke et al., 2023)
and real video data from the MvImgNet (Yu et al., 2023) in training. Results show that removing real
data causes an obvious drop for all the metrics, despite the fact our synthetic 3D dataset contains 3ˆ
more shapes than MvImgNet. One potential reason is that the real data have much more variation
in the lighting, the size of the target, and the camera poses, which effectively benefits the learning.
Future work could augment the rendering of synthetic shapes to adequately utilize those abundant
data. Nevertheless, combining the two datasets leads to substantially better results than training on
any one of them alone.
Table 3: Influence of training datasets.
Data
Unseen Evaluation
PSNRÒ
CLIP-SimilarityÒ
SSIMÒ
LPIPSÓ
Synthetic (Objaverse)
15.5
84.7
70.3
29.3
Real (MvImgNet)
17.5
85.7
75.7
22.0
Synthetic+Real
19.0
87.8
77.4
19.1
D.2
NUMBER OF VIEWS IN TRAINING DATA
In Table 4, we conduct experiments with all data but limit the number of training views per shape.
For example, for Train Views“8, we use only a random subset of 8 views per shape and keep
randomly sampling 4 views from the above subset at each training step. The results show that more
views can lead to better results, possibly because of more diverse data. While the growth is saturated
at 16 views, adding more views does not lead to worse results.
D.3
MODEL HYPER-PARAMETERS
Table 5 presents the results of having a different number of cross-attention layers in the image-to-
triplane decoder. There is a slight trend indicating that the scores can be improved by having a
21
Table 4: Effect of the number of different views per shape in training. 32+ indicates some video
data in MvImgNet contain more than 32 views per shape, which we apply all of them in training.
Train Views
Unseen Evaluation
PSNRÒ
CLIP-SimilarityÒ
SSIMÒ
LPIPSÓ
4
18.8
86.7
77.5
19.8
8
18.9
87.3
77.5
19.4
16
19.1
87.9
77.6
19.0
32+
19.0
87.8
77.4
19.1
deeper model, especially for the latent semantic and perceptual similarity measurements CLIP and
LPIPS, implying that the network models better representations for reconstructing higher-quality
images.
We also evaluate the influence of the number of MLP layers in NeRF (Table 6). Results show that
it is unnecessary to have a very large network, and there seems to be a sweet spot around two to
four layers. This observation is consistent with EG3D (Chan et al., 2022) where the information
of shapes is encoded by the triplane and such MLP is only a shallow model for projecting triplane
features to color and density.
As shown in Table 7, we found that increasing the triplane resolution leads to better image quality.
Note that, in this experiment, we only use a deconvolution layer to upsample the 32ˆ32ˆ32 triplane
produced by LRM’s decoder, whereas we suspect a large improvement could be seen by increasing
the quantity of input spatial-positional embeddings to query more fine-grained image details. How-
ever, such an approach will dramatically increase the computational cost, we leave this exploration
to future research.
Table 5: Effect of the number of cross-attention layers in image-to-triplane decoder.
CrossAttn
Layers
Unseen Evaluation
PSNRÒ
CLIP-SimilarityÒ
SSIMÒ
LPIPSÓ
6
19.0
87.7
77.6
19.1
16
19.0
87.8
77.4
19.1
24
19.1
88.0
77.6
18.9
Table 6: Effect of the number of MLP layers in NeRF.
NeRF MLP
Layers
Unseen Evaluation
PSNRÒ
CLIP-SimilarityÒ
SSIMÒ
LPIPSÓ
2
19.2
87.7
77.8
18.9
6
19.1
88.0
77.6
19.0
12
19.0
87.8
77.4
19.1
14
19.1
87.2
77.6
19.0
Table 7: Effect of the resolution of triplane. For 64up and 128up, we apply additional 2ˆ2 and 4ˆ4
deconvolution layers, respectively, to upsample a Res. 32 triplane.
Triplane Res.
Unseen Evaluation
PSNRÒ
CLIP-SimilarityÒ
SSIMÒ
LPIPSÓ
32
18.9
86.3
77.2
19.7
64up
19.0
87.8
77.4
19.1
128up
19.0
88.3
77.5
19.0
D.4
CAMERA POSE
As we have discussed in the Main Paper, normalizing camera poses in training has a huge impact
on the generalization of input views. We can see from Table 8 that when no modification is ap-
plied (None), LRM produces the worst results. Augmenting camera poses with a Random rotation
22
greatly improves the results since the model learns a more general image-to-triplane projection via
decoupled views and camera poses. However, such unconstrained projection is very difficult to
learn. We therefore Normalized all camera poses so that all images are projected onto the triplane
from the same direction, allowing the model to adequately learn and utilize the cross-shape prior for
reconstruction.
Table 8: Effect of camera pose normalization.
Camera Pose
Unseen Evaluation
PSNRÒ
CLIP-SimilarityÒ
SSIMÒ
LPIPSÓ
None
15.3
83.4
70.1
28.9
Random
18.0
85.6
75.7
21.1
Normalized
19.0
87.8
77.4
19.1
D.5
IMAGE QUANTITY AND RESOLUTION
Table 9 and Table 10 study the influence of the number of side views supervision for each sample
and the effect of image rendering resolution in training. Results indicate that as the quantity of side
views increases, the reconstructed image quality improves. Having more views allows the model
to better correlate the appearance and geometry of different parts of the same shape, and facilitates
inferring multi-view consistent results. Moreover, using a higher rendering resolution of images
in training largely improves the results, as the model is encouraged to learn more high-frequency
details.
Table 9: Influence of the number of side views applied for each training sample.
Side Views
Unseen Evaluation
PSNRÒ
CLIP-SimilarityÒ
SSIMÒ
LPIPSÓ
1
18.7
87.7
77.2
19.7
2
18.7
87.5
77.2
19.6
3
19.0
87.8
77.4
19.1
4
19.1
87.8
77.6
18.9
Table 10: Influence of the rendering resolution of images in training.
Render Res.
Unseen Evaluation
PSNRÒ
CLIP-SimilarityÒ
SSIMÒ
LPIPSÓ
32
18.8
86.3
77.0
20.1
64
19.0
87.8
77.4
19.1
128
19.4
89.0
78.3
18.0
D.6
LPIPS LOSS
Lastly, we found that our LPIPS objective (Zhang et al., 2018) has a huge impact on the results.
Removing it from training will decrease the CLIP-Similarity, SSIM, and LPIPS scores to 74.7, 76.4,
and 29.4, respectively.
E
VISUALIZATIONS
We present more visualizations of the reconstructed 3D shapes in the following pages. The in-
put images include photos captured by our phone camera, images from Objaverse (Deitke et al.,
2023), MvImgNet (Yu et al., 2023), ImageNet (Deng et al., 2009), Google Scanned Objects (Downs
et al., 2022), Amazon Berkeley Objects (Collins et al., 2022), and images generated by the Adobe
Firefly11. We implement a heuristic function to pre-process the camera-captured images, generated
11Adobe Firefly, a text-to-image generation tool: https://firefly.adobe.com.
23
images, and images from MvImgNet and ImageNet. The function removes the image background
with an off-the-shelf package12, followed by cropping out the target object, rescaling the target to a
suitable size and centering the target on a square white figure. All input images are never seen by
the model in training. Please visit our project webpage https://yiconghong.me/LRM/ for video
demonstrations and interactable 3D meshes.
Google Scanned
Objects
Input Image
Rendered
GT
Rendered
GT
Amazon Berkeley
Objects
Objaverse
MvImgNet
Input Image
Rendered
GT
Rendered
GT
MvImgNet
Google Scanned
Objects
Amazon Berkeley
Objects
Figure 6: Comparison between LRM rendered novel views and the ground truth images (GT). None
of the images are observed by the model during training. The GT depth images of Objaverse are
rendered from the 3D models. Please zoom in for clearer visualization.
12Rembg package, a tool to remove image background: https://pypi.org/project/rembg
24
Phone
Captured
Phone
Captured
Phone
Captured
Phone
Captured
ImageNet
ImageNet
ImageNet
ImageNet
Input Image
Rendered Novel Views
Input Image
Rendered Novel Views
ImageNet
ImageNet
Generated
Generated
Generated
Generated
Figure 7: Rendered novel views (RGB and Depth) of shapes reconstructed by our LRM from single
images. None of the images are observed by the model during training. Generated images are
created by the Adobe Firefly. Please zoom in for clearer visualization.
25
Under review as a conference paper at ICLR 2024
INSTANT3D: FAST TEXT-TO-3D WITH SPARSE-VIEW
GENERATION AND LARGE RECONSTRUCTION MODEL
Anonymous authors
Paper under double-blind review
ABSTRACT
Text-to-3D with diffusion models have achieved remarkable progress in recent
years. However, existing methods either rely on score distillation-based optimiza-
tion which suffer from slow inference, low diversity and Janus problems, or are
feed-forward methods that generate low quality results due to the scarcity of 3D
training data. In this paper, we propose Instant3D, a novel method that generates
high-quality and diverse 3D assets from text prompts in a feed-forward manner.
We adopt a two-stage paradigm, which first generates a sparse set of four struc-
tured and consistent views from text in one shot with a fine-tuned 2D text-to-image
diffusion model, and then directly regresses the NeRF from the generated images
with a novel transformer-based sparse-view reconstructor. Through extensive ex-
periments, we demonstrate that our method can generate high-quality, diverse and
Janus-free 3D assets within 20 seconds, which is two order of magnitude faster
than previous optimization-based methods that can take 1 to 10 hours. Our project
webpage: https://instant-3d.github.io/.
1
INTRODUCTION
In recent years, remarkable progress has been achieved in the field of 2D image generation. This
success can be attributed to two key factors: the development of novel generative models such as
diffusion models (Song et al., 2021; Ho et al., 2020; Ramesh et al., 2022; Rombach et al., 2021),
and the availability of large-scale datasets like Laion5B (Schuhmann et al., 2022). Transferring
this success in 2D image generation to 3D presents challenges, mainly due to the scarcity of avail-
able 3D training data. While Laion5B has 5 billion text-image pairs, Objaverse-XL (Deitke et al.,
2023a), the largest public 3D dataset, contains only 10 million 3D assets with less diversity and
poorer annotations. As a result, previous attempts to directly train 3D diffusion models on existing
3D datasets (Luo & Hu, 2021; Nichol et al., 2022; Jun & Nichol, 2023; Gupta et al., 2023; Chen
et al., 2023b) are limited in the visual (shape and appearance) quality, diversity and compositional
complexity of the results they can produce.
To address this, another line of methods (Poole et al., 2022; Wang et al., 2023a; Lin et al., 2023;
Wang et al., 2023b; Chen et al., 2023c) leverage the semantic understanding and high-quality gen-
eration capabilities of pretrained 2D diffusion models. Here, 2D generators are used to calculate
gradients on rendered images, which are then used to optimize for a 3D representation, usually a
NeRF (Mildenhall et al., 2020). Although these methods yield better visual quality and text-3D
alignment, they can be incredibly time-consuming, taking hours for optimizing 3D for each prompt.
They also suffer from artifacts such as over-saturated colors and the “multi-face” problem arising
from the bias in pretrained 2D diffusion models, and struggle to generate diverse results from the
same text prompt, with varying the random seed leading to minor changes in geometry and texture.
In this paper, we propose Instant3D, a novel feed-forward method that generates high-quality and
diverse 3D assets conditioned on text. Instant3D, like the methods noted above, builds on top of
pretrained 2D diffusion models. However, it does so by splitting 3D generation into two stages:
2D generation and 3D reconstruction. In the first stage, instead of generating images sequentially,
we fine-tune an existing text-to-image diffusion model (Podell et al., 2023) to generate a sparse set
of four-view images in the form of a 2×2 grid in a single denoising process. This design allows
the multi-view images to attend to each other during generation, leading to more view-consistent
results. In the second stage, instead of relying on a slow optimization-based reconstruction method,
1
Under review as a conference paper at ICLR 2024
a steaming basket full of dumplings
a steam engine train, high resolution
a 3d model of an adorable cottage with a thatched roof
a baby dragon drinking boba
beautifully carved jade heavenly Chinese guardian lion,
hand carved nephrite jade
a ghost eating a hamburger
a hippo wearing a sweater
a panda rowing a boat in a pond
a snail on a leaf
a cracked egg with the yolk spilling out
Figure 1: Our method generates high-quality 3D NeRF assets from the given text prompts within 20
seconds. Here we show novel view renderings from our generated NeRFs as well as the renderings
of the extracted meshes from their density field.
inspired by (Anonymous, 2023), we introduce a novel sparse-view large reconstruction model with
a transformer-based architecture that can directly regress a triplane-based (Chan et al., 2022) NeRF
from a sparse set of multi-view images. Our model projects sparse-view images into a set of pose-
aware image tokens using pretrained vision transformers (Caron et al., 2021), which are then fed
to an image-to-triplane decoder that contains a sequence of transformer blocks with cross-attention
and self-attention layers. Our proposed model has a large capacity with more than 500 million
parameters and can robustly infer correct geometry and appearance of objects from just four images.
Both of these stages are fine-tuned/trained with multi-view rendered images of around 750K 3D
objects (Deitke et al., 2023b). Both of these stages are fine-tuned/trained with multi-view rendered
images of 3D objects (Deitke et al., 2023b). The second reconstruction stage takes around 750K
multi-view data in training while the first stage can take as little as 10K data. While we use rel-
atively smaller dataset compared to the pre-training dataset for other modalities (e.g., C4 Raffel
et al. (2020) for text and Laion5B for image), Instant3D’s two-stage approach combines it with the
power of pretrained 2D diffusion models to generate high-quality, diverse 3D assets even for com-
plex input prompts (see Figure 1). Due to its feed-forward architecture, Instant3D is exceptionally
fast, requiring only about 20 seconds for a text prompt, which is about 200× faster than previous
optimization-based methods (Poole et al., 2022; Wang et al., 2023b) while achieving comparable or
even better quality.
2
RELATED WORKS
2.1
3D GENERATION
Following the success of generative models on 2D images using GANs (Goodfellow et al., 2014;
Karras et al., 2019; Gu et al., 2022; Kang et al., 2023) and diffusion models (Rombach et al., 2021;
Podell et al., 2023; Ho et al., 2020; Song et al., 2021; Saharia et al., 2022), people have also explored
the applications of such models on 3D generation. Previously GAN-based approached have explored
different methods to generate 3D models in the form of point clouds (Wu et al., 2016; Gadelha
et al., 2017; Smith & Meger, 2017), triangle meshes (Gao et al., 2022; Pavllo et al., 2020; Chen
et al., 2019; Luo et al., 2021) or volumes (Chan et al., 2022; Or-El et al., 2022; Bergman et al.,
2022; Skorokhodov et al., 2022) in either unconditional or text/image-conditioned manner, typically
supervised by adversarial losses on the 3D representations or 2D rendered images. Such methods
usually require category-specific training and do not generalize well to novel categories.
2
Under review as a conference paper at ICLR 2024
Diffusion models open new possibilities for 3D generation. A class of methods directly train 3D
diffusion models on the 3D representations (Nichol et al., 2022; Liu et al., 2023c; Zhou et al., 2021)
or project the 3D models or multi-view rendered images into latent representations (Ntavelis et al.,
2023; Zeng et al., 2022; Gupta et al., 2023; Jun & Nichol, 2023; Chen et al., 2023b) and perform the
diffusion process in the latent space. For example, Shap-E (Jun & Nichol, 2023) encodes each 3D
shape into a set of parameters of an implicit function, and then trains an conditional diffusion models
on the parameters. These approaches face challenges due to the restricted availability and diversity of
existing 3D data, consequently resulting in generated content with poor visual quality and inadequate
lignment with the input prompt. Therefore, although trained on millions of 3D assets, Shap-E still
fails to generate 3D shapes with complex compositional concepts and high-fidelity textures.
To resolve this, another line of works try to make use of 2D diffusion models to facilitate the 3D gen-
eration. Some works (Jain et al., 2022; Mohammad Khalid et al., 2022) optimize meshes or NeRFs
to minimize the CLIP Radford et al. (2021) score between the rendered images and input prompt
utilizing pretrained CLIP model. While such methods can generate diverse 3D content, they exhibit
a deficiency in achieving visual realism. More recently, some works (Poole et al., 2022; Wang et al.,
2023b; Lin et al., 2023; Chen et al., 2023c) optimize 3D representations using score distillition loss
(SDS) based on a pretrained 2D diffusion models. Such methods can generate high-quality results,
but suffer from the slow optimization, over-saturated colors and the Janus problem. For example, it
takes 1.5 hours for DreamFusion (Poole et al., 2022) and 10 hours for ProlificDreamer Wang et al.
(2023b) to generate a single 3D asset, which greatly limits their practicality. In contrast, our method
enjoys the benefits of both worlds: it’s able to borrow information from pretrained 2D diffusion
models to generate diverse multi-view consistent images that are subsequently lifted to faithful 3D
models, while still being fast and efficient due to its feed-forward nature.
2.2
SPARSE-VIEW RECONSTRUCTION
Traditional 3D reconstruction with multi-view stereo (Agarwal et al., 2011; Sch¨
onberger et al., 2016;
Furukawa et al., 2015) typically require a dense set of input images that have significant overlaps
to find correspondence across views and infer the geometry correctly. While NeRF (Mildenhall
et al., 2020) and its variants (M¨
uller et al., 2022; Chen et al., 2022; 2023a) have further alleviated
the prerequisites for 3D reconstruction, they perform per-scene optimization that still necessitate a
lot of input images. Previous methods (Wang et al., 2021; Chen et al., 2021; Long et al., 2022;
Reizenstein et al., 2021; Trevithick & Yang, 2021) have tried to learn data priors so as to infer NeRF
from a sparse set of images. Typically they extract per-view features from each input image, and then
for each point on the camera ray, aggregate the multi-view features and decode them to the density
(or SDF) and colors. Such methods are either trained in a category-specific manner, or only trained
on a small datasets such as ShapeNet and ScanNet; they have not been demonstrated to generalize
beyond these datasets especially to the complexity of text-to-2D outputs.
More recently, some methods utilize data priors from pretrained 2D diffusion models to lift a single
2D image to 3D by providing supervision at novel views using SDS loss (Liu et al., 2023b; Qian
et al., 2023; Melas-Kyriazi et al., 2023) or generated multi-view images (Liu et al., 2023a). For
instance, One-2-3-45 Liu et al. (2023a) generates 32 images at novel views from a single input image
using a fine-tuned 2D diffusion model, and reconstructs a 3D model from them; this leads to issues
with inconsistency between the many input views. In comparison, our sparse-view reconstructor
adopts a highly scalable transformer-based architecture and is trained on large-scale 3D data. This
gives it the ability to accurately reconstruct 3D models of novel unseeen objects from a sparse set
of 4 images without per-scene optimization. In a way, our work can be seen as revisiting the earlier
attempts at 3D reconstruction from sparse views, such as space carving, visual hulls and photo
hulls (Kutulakos & Seitz, 2000), but with the benefit of modern learning techniques, and in the
context of generation rather than purely scene reconstruction.
3
METHOD
Our method Instant3D is composed of two stages: sparse-view generation and feed-forward NeRF
reconstruction. In Section 3.1, we present our approach on generating sparse multi-view images
conditioned on the text input. In Section 3.2, we describe our transformer-based sparse-view large
reconstruction model.
3
Under review as a conference paper at ICLR 2024
A car made
out of sushi
2D diffusion model
(fine-tuned)
Transformer-based
reconstructor
Gaussian blob initialization
(inference only)
4-view image (2 by 2 grid)
NeRF
Figure 2: Overview of our method. Given a text prompt (‘a car made out of sushi’), we perform
multi-view generation with Gaussian blobs as initialization using fine-tuned 2D diffusion model,
producing a 4-view image in the form of a 2 by 2 grid. Then we apply a transformer-based sparse-
view 3D reconstructor on the 4-view image to generate the final NeRF.
3.1
TEXT-CONDITIONED SPARSE VIEW GENERATION
Given a text prompt, our goal is to generate a set of multi-view images that are aligned with the
prompt and consistent with each other. We achieve this by fine-tuning a pretrained text-to-image
diffusion models to generate 2x2 grid as shown in Figure 2.
In the following paragraphs, we first illustrate that the large text-to-image diffusion model (i.e.,
SDXL (Podell et al., 2023)) has the capacity to generate view-consistent images thus a light-weight
fine-tuning is possible. We then introduces three essential techniques to realize it: the image grid,
the curation of the dataset, and also the Gaussian Blob noise initialization in inference. As a result
of these observations and technical improvements, we can fine-tune the 2D diffusion model for only
10K steps (on 10K data) to generate consistent sparse views.
Multi-view generation with image grid. Previous methods (Liu et al., 2023b;a) on novel-view
synthesis show that image diffusion models are capable of understanding the multi-view consistency.
In light of this, we compile the images at different views into a single image in the form of an image
grid, as depicted in Figure 2. This image-grid design can better match the original data format of
the 2D diffusion model, and is suitable for simple direct fine-tuning protocol of 2D models. We
also observe that this simple protocol only works when the base 2D diffusion has enough capacity,
as shown in the comparisons of Stable Diffusion v1.5 (Rombach et al., 2021) and SDXL (Podell
et al., 2023) in Section 4.3. The benefit from simplicity will also be illustrated later in unlocking the
light-weight fine-tuning possibility.
Regarding the number of views in the image grid, there is a trade-off between the requirements of
the multi-view generation and 3D reconstructor. More generated views make the problem of 3D re-
construction easier with more overlaps but increase possibility of view inconsistencies in generation.
On the other hand, too few views may cause insufficient coverage, requiring the reconstructor to hal-
lucinate unseen parts, which is challenging for a deterministic 3D reconstruction model. Thanks to
our sparse-view constructor which greatly reduces the number of views, we empirically found that
4 views is a balance between these two needs. Moreover, 4 views can be naturally arranged in a 2x2
grid as detailed in Figure 2. Next, we detail how the image grid data is created and curated.
Multi-view data creation and curation.
To fine-tune the text-to-image diffusion model, we create
paired multi-view renderings and text prompts. We adopt a large-scale synthetic 3D dataset (Deitke
et al., 2023b) and render four 512x512 views of about 750K objects with Blender. We distribute the
four views at a fixed elevation (20 degrees) and four equidistant azimuths (0, 90, 180, 270 degrees)
to achieve a better coverage of the object. We use the captions from Cap3D (Luo et al., 2023) ,
which utilizes an off-the-shelf image captioning model BLIP-2 (Li et al., 2023). As the final pre-
processing step, the four view are then assembled into a grid image in a fixed order and resized to
the input resolution compatible of the 2D diffusion model.
We find that naively using all the data for fine-tuning has a negative impact on the quality of the
results. Therefore, we train a simple scorer on a small amount (2000 samples) of manually labelled
data to predict asset quality. The model is a simple SVM on top of pretrained CLIP features extracted
from multi-view renderings of a 3D object. See Appendix for details.
4
Under review as a conference paper at ICLR 2024
Self-
attention
MLP
+
+
12 layers, ViT (DINO)
Camera features
Conv
Conv
Multi-view images
(N=2 here)
Image encoder
(shared)
Image encoder
Triplane tokens
(learnable)
Concat
2D image tokens
Self-
attention
MLP
+
16 layers
Cross-
attention
reshape &
upsample
Triplane
Volume
rendering
+
+
Mod
Novel view rendering
Image-to-triplane decoder
Figure 3: Architecture of our sparse-view reconstructor. The model applies a pretrained ViT to
encode multi-view images into pose-aware image tokens, from which we decode a triplane rep-
resentation of the scene using a transformer-based decoder. Finally we decode per-point triplane
features to its density and color and perform volume rendering to render novel views. We illustrate
here with 2 views and the actual implementation uses 4 views.
Our model only takes the top 10K data ranked by our scorer for training by default. We provide
quantitative study in Section 4.3 for the impact of more or less data. We also provide the results
from our uncurated data. Although the difference is not big from the metric side, we found that this
curated data is helpful in quality.
Inference with Gaussian blob initialization.
While our data are multi-view images with white
background, we observed that during inference starting from standard Gaussian noise results in low-
quality images that have cluttered backgrounds; this introduces extra difficulty for the feed-forward
reconstructor in the second stage (Section 3.2). To guide the model toward generating images with
a clean white backgrounds, inspired by SDEdit (Meng et al., 2022), we first create a white image
of a 2x2 grid that has the same resolution as the output images, and initialize each sub-grid with
a black 2D Gaussian blob that is placed at the center of image with a standard deviation of 0.1.
The visualization of this Gaussian Blob can be found in Figure 2; please refer to Appendix for the
concrete math formulation. This Gaussian blob image grid is fed to the autoencoder to get its the
latent representation. We then add diffusion noise (e.g., use t=980/1000 for 50 DDIM denoising
steps), and use it as the starting point for the denoising process. As seen in Figure 5, this technique
effectively guides the model toward generating images with clean background.
Light-weight fine-tuning.
With all above observations and techniques, we are able to adapt a
text-to-image diffusion model to a text-to-multiviews model with light-weight fine-tuning. This
light-weight fine-tuning share a similar spirit as the ‘instruction fine-tuning’ (Mishra et al., 2022;
Wei et al., 2021) for LLM (large language model) alignment. The assumption is that the base model
is already capable of the task, and the fine-tuning is to unlock the base model’s ability without
introducing additional knowledge.
Since we utilize image grid, the fine-tuning follows the exact same training protocol as the 2D
diffusion model pre-training, except that we decrease the learning rate to 10−5. We train the model
with a batch size of 192 for only 10K iterations on the 10K curated multi-view data. The training is
done on 32 NVIDIA A100 GPUs for only 3 hours. We study the impact of different training amount
in Section 4.3. For more training details, please refer to Appendix.
3.2
FEED-FORWARD SPARSE-VIEW LARGE RECONSTRUCTION MODEL
In this stage, we aim to reconstruct a NeRF from the four-view images I = {Ii | i = 1, ..., 4}
generated in the first stage. 3D reconstruction from sparse inputs with a large baseline is a challeng-
ing problem, which requires strong model priors to resolve the inherent ambiguity. Inspired by a
recent work LRM (Anonymous, 2023) that introduces a transformer-based model for single image
5
Under review as a conference paper at ICLR 2024
Table 1: Quantitative comparisons on CLIP scores
against baseline methods.
Our method outper-
forms previous feed-forward method Shap-E and
optimization-based
method
DreamFusion,
and
achieves competitive performance to ProlificDreamer
while being 1800× faster.
ViT-L/14 ↑
ViT-bigG-14 ↑
Time(s) ↓
Shap-E
20.51
32.21
6
DreamFusion
23.60
37.46
5400
ProlificDreamer
27.39
42.98
36000
Ours
26.87
41.77
20
Table 2:
Quantitative comparisons
against previous sparse-view recon-
struction methods on GSO dataset.
PSNR ↑
SSIM ↑
LPIPS ↓
SparseNeus
20.62
0.8360
0.1989
Ours
26.54
0.8934
0.0643
3D reconstruction, we propose a novel approach that enables us to predict a NeRF from a sparse set
of input views with known poses. Similar to Anonymous (2023), our model consists of an image
encoder, a images-to-triplane decoder, and a NeRF decoder. The image encoder encodes the multi-
view images into a set of tokens. We feed the concatenated image tokens to the image-to-triplane
decoder to output a triplane representation (Chan et al., 2022) for the 3D object. Finally, the triplane
features are decoded into per-point density and colors via the NeRF MLP decoder.
In detail, we apply a pretrained Vision Transformer (ViT) DINO (Caron et al., 2021) as our image
encoder. To support multi-view inputs, we inject camera information in the image-encoder to make
the output image tokens pose-aware. This is different from Anonymous (2023) that feeds the camera
information in the image-to-triplane decoder because they take single-image input. The camera
information injection is done by the AdaLN (Huang & Belongie, 2017; Peebles & Xie, 2022) camera
modulation as described in Anonymous (2023). The final output of the image encoder is a set of
pose-aware image features f ∗
Ii, and we concatenate the per-view features together as the feature
descriptors for the multi-view images: fI = ⊕(f ∗
I1, ...f ∗
I4)
We use triplane as the scene representation. The triplane is flattened to a sequence of learnable to-
kens, and the image-to-triplane decoder connects these triplane tokens with the pose-aware image
tokens fI using cross-attention layer, followed by self-attention and MLP layers. The final out-
put tokens are reshaped and upsampled using a de-convolution layer to the final triplane features.
During training, we ray-march through the object bounding box and decode the triplane features at
each point to its density and color using a shared MLP, and finally get the pixel color via volume
rendering. We train the networks in an end-to-end manner with image reconstruction loss at novel
views using a combination of MSE loss and LPIPS (Zhang et al., 2018) loss.
Training Details.
We train the model on multi-view renderings of the Ojbaverse dataset (Deitke
et al., 2023b). Different from the first stage that performs data curation, we use all the 3D data in
the dataset and scale them to [−1, 1]3; then we generate multi-view renderings using Blender under
uniform lighting with a resolution of 512 × 512. While the output images from the first stage is
generated in a structured setup with fixed camera poses, we train the model using random views as a
data augmentation mechanism to increase the robustness. Particularly, for each object, we randomly
sample 32 views around the object. During training, for each object we randomly select a subset
of 4 images as input and another random set of 4 images as supervision. During inference, we will
reuse the fixed camera poses in the first stage as the camera input to the second stage. For more
details on the training, please refer to the Appendix.
4
EXPERIMENTS
In this section, we first do comparisons against previous methods on text-to-3D (Section 4.1), and
then perform ablation studies on different design choices of our method. By default, we report the
results generated with fine-tuned SDXL models, unless otherwise noted.
4.1
TEXT-TO-3D
We make comparisons to state-of-the-art methods on text-to-3D, including feed-forward methods
Shap-E (Jun & Nichol, 2023), and optimization-based methods including DreamFusion (Poole et al.,
6
Under review as a conference paper at ICLR 2024
Figure 4: Qualitative comparisons on text-to-3D compared against previous methods. We include
more uncurated comparison results in the supplementary material.
2022) and ProlificDreamer (Wang et al., 2023b). We use the official code for Shap-E, and use the
implementation from three-studio (Guo et al., 2023) for the other two as there is no official code.
We keep all the hyper-parameters (number of optimization iterations, number of denoising steps) of
these models as default. For our own model we use the SDXL base model trained on 10K data for
10K steps. During inference we take 100 DDIM steps.
Qualitative comparisons.
As shown in Figure 4, our method generates visually better results
than those of Shap-E, producing sharper textures, better geometry and substantially improved text-
3D alignment. Shap-E applies a 3D diffusion model that are exclusively trained on million-level
3D data, which might be an evidence for the need of 2D data or 2D prior models. DreamFusion
and ProlificDreamer achieves better text-3D alignment utilizing a pretrained 2D diffusion model.
However, DreamFusion generates results with over-saturated colors and over-smooth textures. Pro-
lificDreamer results have better details. However, it still suffer from low-quality geometry (as in ‘A
bulldozer clearing ...’) and the Janus problem (as in ”a squirrel dressed like ...”, also more detailed
in Appendix Figure 11). In comparison, our results have more photorealistic colors with higher
geometric details. Please refer to the Appendix and supplementary materials for video comparisons
and more results.
Quantitative comparisons.
In Table 4, we quantitatively assess the coherence between the gen-
erated models and the text prompts using CLIP-based scores. We perform the evaluation on re-
sults with 400 text prompts from DreamFusion. For each model, we render 10 random views and
calculate the average CLIP score between the rendered images and the input text. We report the
metric using multiple variants of CLIP models with different model size and training data (i.e., ViT-
L/14 from OpenAI and ViT-bigG-14 from OpenCLIP). From the result we can see that our model
achieves higher CLIP scores than Shap-E, indicating better text-3D alignment. Our method even
achieves consistently higher CLIP scores than optimization-based method DreamFusion and com-
petitive scores to ProlificDreamer, from which we can see that our approach can effectively inherit
the great text understanding capability from pretrained SDXL models while preserve them in the
generated 3D assets via consistent sparse-view generation and robust 3D reconstruction.
Inference time comparisons.
We present the time to generate one 3D assets in Table 1. The
timing is calculated using the default hyper-parameters of each method on an A100 GPU. Notably,
7
Under review as a conference paper at ICLR 2024
a barbecue grill cooking sausages and burger patties 001
Frog 004
w/o Gaussian blobs
w/ Gaussian blobs
Original
Finetuned
Figure 5: Qualitative comparisons on results generated with and without Gaussian blob initialization.
our method is significantly faster than the optimization-based methods: while it takes 1.5 hours for
DreamFusion and 10 hours for ProlificDreamer to generate a single asset, our method can finish
the generation within 20 seconds, resulting in a 270× and 1800× speed up respectively. In Fig-
ure 10, we show that our inference time can be further reduced by reducing the DDIM steps without
significantly sacrificing the quality.
4.2
COMPARISONS ON SPARSE VIEW RECONSTRUCTION
We make comparison to previous sparse-view NeRF reconstruction works.
Most of previous
works (Reizenstein et al., 2021; Trevithick & Yang, 2021; Yu et al., 2021) are either trained on
a small-scale dataset such as ShapeNet, or trained in a category-specific manner. Therefore, we
make comparison to a state-of-the-art method SparseNeus (Long et al., 2022), which is also applied
in One2345 (Liu et al., 2023a) where they train the model on the same Objaverse dataset for sparse-
view reconstruction. We do the comparisons on the Google Scan Object (GSO) dataset (Downs
et al., 2022), which consists of 1019 objects. For each object, we render 4-view input following the
structured setup and randomly select another 10 views for testing. We adopt the pretrained model
from Liu et al. (2023a). Particularly, SparseNeus does not work well for 4-view inputs with such a
large baseline; therefore we add another set of 4 input views in addition to the our four input views,
following the setup in Liu et al. (2023a). We report the metrics on novel view renderings in Table 2.
For the table, we can see that our method outperforms the baseline method even with fewer input
images, which demonstrates the superiority of our sparse-view reconstructor.
4.3
ABLATION STUDY FOR SPARSE VIEW GENERATION
We ablate several key decisions in our method design, i.e., (1) the choice of larger 2D base model
SDXL, (2) the use of Guassian Blob during inference, (3) the quality and size of curated dataset,
also lastly, (4) the need and requirements of light-weight fine-tuning. We gather the quantitative
results in Table 3 and put all qualitative results in Appendix. We observe that qualitative results are
more evident than quantitative results thus recommend a look.
Scalability with 2D text-to-image models.
One of the notable advantages of our method is that
the efficacy of our method scales positively with the potency of the underlying 2D text-to-image
model. In Figure 12, we present qualitative comparisons between two distinct backbones (with
their own tuned hyperparameters): SD1.5 (Rombach et al., 2021) and SDXL (Podell et al., 2023).
It becomes readily apparent that SDXL, which boasts a model size 3× larger than that of SD1.5,
exhibits superior text comprehension and visual quality. We also show a quantitative comparison on
CLIP scores in Table 3, by comparing with Exp(l, m) with Exp(d, g), we can see that the model with
SD1.5 achieves consistently lower CLIP scores indicating worse text alignment.
Gaussian Blobs Initialization.
In Figure 5, we show our results generated with and without Gaus-
sian blob initialization. From the results we can see that while our fine-tuned model can generate
multi-view images without Gaussian blob initialization, they tend to have cluttered backgrounds,
which challenges the second-stage feed-forward reconstructor. In contrast, our proposed Gaussian
blob initialization enables the fine-tuned model to generate images with clean white background,
which better align with the image input data in the second stage.
8
Under review as a conference paper at ICLR 2024
Table 3: Comparison on CLIP scores of NeRF renderings with different variants of fine-tuning
settings.
Exp ID
Exp Name
Base
# Data
Curated
# Steps
ViT-L/14
ViT-bigG-14
(a)
Curated-1K-s1k
SDXL
1K
✓
1K
26.33
41.09
(b)
Curated-1K-s10k
SDXL
1K
✓
10k
22.55
35.59
(c)
Curated-10K-s4k
SDXL
10K
✓
4k
26.55
41.08
(d)
Curated-10K-s10k
SDXL
10K
✓
10k
26.87
41.77
(e)
Curated-10K-s20k
SDXL
10K
✓
20k
25.96
40.56
(f)
Curated-100K-s10k
SDXL
100K
✓
10k
25.79
40.32
(g)
Curated-100K-s40k
SDXL
100K
✓
40k
26.59
41.29
(h)
Curated-300K-s40k
SDXL
300K
✓
40K
26.43
40.72
(i)
Random-10K-s10k
SDXL
10K
✗
10k
26.87
41.47
(j)
Random-100K-s40k
SDXL
100K
✗
40k
26.28
40.90
(k)
AllData-s40k
SDXL
700K
✗
40k
26.13
40.60
(l)
Curated-10K-s10k (SD1.5)
SD1.5
10K
✓
10k
23.50
36.90
(m)
Curated-100K-s40k (SD1.5)
SD1.5
100K
✓
40k
25.48
39.07
Size and quality of fine-tuning dataset.
We evaluate the impact of the quality and size of the
dataset used for fine-tuning 2D text-to-image models. We first make comparisons between curated
and uncurated randomly selected data. The CLIP alignment score raises slightly as shown in Table 3
(i.e., comparing Exp ID d&i), while there is a substantial quality improvement as illustrated in
Appendix Figure 7. This aligns with the observation that the data quality can dramatically impact
the results in the instruction fine-tuning stage for LLM (Zhou et al., 2023).
When it comes to data size, we observe a double descent from Table 3 Exp ID (a, d, g) with 1K, 10K,
and 100K data. We pick experiment (a, d, g) here because they are the best results among different
training steps for the same training data size. The reason for this double descent can be spotlighted
by the qualitatively comparisons in Appendix Figure 13, where training with 1K data can lead to
incomplete and inconsistent images, while training with 100K data can hurt the compositionality,
photo realism, and also the text alignment.
Number of fine-tuning steps.
We also both quantitatively and qualitatively analyze the impact
of fine-tuning steps. For each block Table 3 in the Table, we show the CLIP text alignment scores
for different training setups. Similar to the findings in instruction fine-tuning (Ouyang et al., 2022),
the results do not increase monotonically regarding the number of fine-tuning steps but has a peak
in middle. For example, in our final setup trained with SDXL base model and 10K cureated data
(i.e., Exp ID c, d, e), the results are peaked at 10K step and we also take this number to train our
final model. For other data setups and number of data amount, the observations are similar. We
also qualitatively visualize the result comparisons at different training steps for 10K curated data in
Appendix Figure 14. There is observable result degradation for both 4K and 20K training steps.
Another important observation is that the peak might move earlier when the model size goes larger.
This can be observed by comparing between SD1.5 (Exp ID l&m in Table 3) and SDXL (Exp ID
d&g)) results. Note that this comparison is not conclusive yet from the Table given that SD1.5 does
not perform reasonbily with our direct fine-tuning protocal. More details are in Appendix.
We also found that Exp ID (a) with 1K steps on 1K data can achieve the best CLIP scores but the
the view consistency is actually disrupted. A possible reason is that the CLIP score is insensitive
to certain noise introduced by reconstruction from inconsistent images, which also calls for a more
reliable 3D generation evaluation metric.
5
CONCLUSIONS
In this paper we presented a novel feed-forward two-stage approach Instant3D that can generate
high-quality and diverse 3D assets from text prompts within 20 seconds. Our method leverages
a fine-tuned 2D text-to-image model to generate consistent 4-view images, and lift them to 3D
with a robust transformer-based large reconstruction model. The experiment results show that our
method outperforms previous feed-forward methods in terms of quality while being equally fast, and
achieves comparable or better performance to previous optimization-based methods with a speed-
up of more than 200× . Instant3D allows novice user to easily create 3D assets and enables fast
prototyping and iteration for various applications such as 3D design and modeling.
9
Under review as a conference paper at ICLR 2024
Ethics Statement.
The generation ability of our model is inherited from the public 2D stable dif-
fusion model (SDXL). We only do light-weight fine-tuning over the SDXL model thus it is hard to
introduce extra knowledge to SDXL. Also, our model can share similar ethical and legal consid-
eration as SDXL. The curation of the data for light-weight fine-tuning does not introduce outside
annotators. Thus the quality of the data might be biased towards the preference of authors, which
can lead to a potential bias on the generated results as well. The text input to the model is not further
checked by the model. That means that the model will try generating for every text prompt it gets,
without has the ability to acknowledge unknown knowledge.
Reproducibility Statement.
In the main text, we highlight the essential techniques to build our
model for both the first stage (Section 3.1) and the second stage (Section 3.2). We also discuss how
our data is created and curated in Section 3. The full model configurations and training details can be
found in Appendix Section A.2 and Section A.4. We have detailed all the optimizer hyperparameters
and model dimensions. We also detail our data curation process in Section A.3. Since different
people might have different preference thus lead to different curated dataset, we also attach our
curated data annotation in Supplementary Materials to enforce reproducibility.
REFERENCES
Adobe. Adobe Firefly. https://firefly.adobe.com/, 2023.
Sameer Agarwal, Yasutaka Furukawa, Noah Snavely, Ian Simon, Brian Curless, Steven M Seitz, and
Richard Szeliski. Building rome in a day. Communications of the ACM, 54(10):105–112, 2011.
Anonymous. Lrm: Large reconstruction model for single image to 3d. In Supplementary Files,
2023.
Alexander Bergman, Petr Kellnhofer, Wang Yifan, Eric Chan, David Lindell, and Gordon Wet-
zstein. Generative neural articulated radiance fields. Advances in Neural Information Processing
Systems, 35:19900–19916, 2022.
Mathilde Caron, Hugo Touvron, Ishan Misra, Herv´
e J´
egou, Julien Mairal, Piotr Bojanowski, and
Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of
the International Conference on Computer Vision (ICCV), 2021.
Eric R. Chan, Connor Z. Lin, Matthew A. Chan, Koki Nagano, Boxiao Pan, Shalini De Mello,
Orazio Gallo, Leonidas Guibas, Jonathan Tremblay, Sameh Khamis, Tero Karras, and Gordon
Wetzstein. Efficient geometry-aware 3D generative adversarial networks. In CVPR, 2022.
Anpei Chen, Zexiang Xu, Fuqiang Zhao, Xiaoshuai Zhang, Fanbo Xiang, Jingyi Yu, and Hao Su.
Mvsnerf: Fast generalizable radiance field reconstruction from multi-view stereo. In Proceedings
of the IEEE/CVF International Conference on Computer Vision, pp. 14124–14133, 2021.
Anpei Chen, Zexiang Xu, Andreas Geiger, Jingyi Yu, and Hao Su. Tensorf: Tensorial radiance
fields. In European Conference on Computer Vision (ECCV), 2022.
Anpei Chen, Zexiang Xu, Xinyue Wei, Siyu Tang, Hao Su, and Andreas Geiger. Dictionary fields:
Learning a neural basis decomposition. ACM Trans. Graph., 2023a.
Hansheng Chen, Jiatao Gu, Anpei Chen, Wei Tian, Zhuowen Tu, Lingjie Liu, and Hao Su. Single-
stage diffusion nerf: A unified approach to 3d generation and reconstruction. In ICCV, 2023b.
Rui Chen, Yongwei Chen, Ningxin Jiao, and Kui Jia. Fantasia3d: Disentangling geometry and
appearance for high-quality text-to-3d content creation. arXiv preprint arXiv:2303.13873, 2023c.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear
memory cost. arXiv preprint arXiv:1604.06174, 2016.
Wenzheng Chen, Huan Ling, Jun Gao, Edward Smith, Jaakko Lehtinen, Alec Jacobson, and Sanja
Fidler. Learning to predict 3d objects with an interpolation-based differentiable renderer. Ad-
vances in neural information processing systems, 32, 2019.
10
Under review as a conference paper at ICLR 2024
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher R´
e. Flashattention: Fast and memory-
efficient exact attention with io-awareness. Advances in Neural Information Processing Systems,
35:16344–16359, 2022.
Matt Deitke, Ruoshi Liu, Matthew Wallingford, Huong Ngo, Oscar Michel, Aditya Kusupati, Alan
Fan, Christian Laforte, Vikram Voleti, Samir Yitzhak Gadre, et al. Objaverse-xl: A universe of
10m+ 3d objects. arXiv preprint arXiv:2307.05663, 2023a.
Matt Deitke, Dustin Schwenk, Jordi Salvador, Luca Weihs, Oscar Michel, Eli VanderBilt, Ludwig
Schmidt, Kiana Ehsani, Aniruddha Kembhavi, and Ali Farhadi. Objaverse: A universe of anno-
tated 3d objects. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pp. 13142–13153, 2023b.
Laura Downs, Anthony Francis, Nate Koenig, Brandon Kinman, Ryan Hickman, Krista Reymann,
Thomas B McHugh, and Vincent Vanhoucke. Google scanned objects: A high-quality dataset
of 3d scanned household items. In 2022 International Conference on Robotics and Automation
(ICRA), pp. 2553–2560. IEEE, 2022.
Yasutaka Furukawa, Carlos Hern´
andez, et al.
Multi-view stereo: A tutorial.
Foundations and
Trends® in Computer Graphics and Vision, 9(1-2):1–148, 2015.
Matheus Gadelha, Subhransu Maji, and Rui Wang. 3d shape induction from 2d views of multiple
objects. In 2017 International Conference on 3D Vision (3DV), pp. 402–411. IEEE, 2017.
Jun Gao, Tianchang Shen, Zian Wang, Wenzheng Chen, Kangxue Yin, Daiqing Li, Or Litany, Zan
Gojcic, and Sanja Fidler. Get3d: A generative model of high quality 3d textured shapes learned
from images. Advances In Neural Information Processing Systems, 35:31841–31854, 2022.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair,
Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information
processing systems, 27, 2014.
Jiatao Gu, Lingjie Liu, Peng Wang, and Christian Theobalt. Stylenerf: A style-based 3d aware
generator for high-resolution image synthesis. In International Conference on Learning Repre-
sentations, 2022.
Yuan-Chen Guo, Ying-Tian Liu, Ruizhi Shao, Christian Laforte, Vikram Voleti, Guan Luo, Chia-
Hao Chen, Zi-Xin Zou, Chen Wang, Yan-Pei Cao, and Song-Hai Zhang. threestudio: A unified
framework for 3d content generation. https://github.com/threestudio-project/
threestudio, 2023.
Anchit Gupta, Wenhan Xiong, Yixin Nie, Ian Jones, and Barlas O˘
guz.
3dgen: Triplane latent
diffusion for textured mesh generation. arXiv preprint arXiv:2303.05371, 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in
neural information processing systems, 33:6840–6851, 2020.
Xun Huang and Serge Belongie. Arbitrary style transfer in real-time with adaptive instance normal-
ization. In Proceedings of the IEEE international conference on computer vision, pp. 1501–1510,
2017.
Ajay Jain, Ben Mildenhall, Jonathan T Barron, Pieter Abbeel, and Ben Poole. Zero-shot text-guided
object generation with dream fields. In Proceedings of the IEEE/CVF Conference on Computer
Vision and Pattern Recognition, pp. 867–876, 2022.
Heewoo Jun and Alex Nichol. Shap-e: Generating conditional 3d implicit functions. arXiv preprint
arXiv:2305.02463, 2023.
Minguk Kang, Jun-Yan Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung
Park. Scaling up gans for text-to-image synthesis. In Proceedings of the IEEE Conference on
Computer Vision and Pattern Recognition (CVPR), 2023.
11
Under review as a conference paper at ICLR 2024
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative
adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern
recognition, pp. 4401–4410, 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint
arXiv:1412.6980, 2014.
Kiriakos N Kutulakos and Steven M Seitz. A theory of shape by space carving. International journal
of computer vision, 38:199–218, 2000.
Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi. Blip-2: Bootstrapping language-image
pre-training with frozen image encoders and large language models. 2023.
Chen-Hsuan Lin, Jun Gao, Luming Tang, Towaki Takikawa, Xiaohui Zeng, Xun Huang, Karsten
Kreis, Sanja Fidler, Ming-Yu Liu, and Tsung-Yi Lin. Magic3d: High-resolution text-to-3d con-
tent creation.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pp. 300–309, 2023.
Minghua Liu, Chao Xu, Haian Jin, Linghao Chen, Mukund Varma T, Zexiang Xu, and Hao Su.
One-2-3-45: Any single image to 3d mesh in 45 seconds without per-shape optimization, 2023a.
Ruoshi Liu, Rundi Wu, Basile Van Hoorick, Pavel Tokmakov, Sergey Zakharov, and Carl Vondrick.
Zero-1-to-3: Zero-shot one image to 3d object, 2023b.
Zhen Liu, Yao Feng, Michael J Black, Derek Nowrouzezahrai, Liam Paull, and Weiyang Liu.
Meshdiffusion: Score-based generative 3d mesh modeling. arXiv preprint arXiv:2303.08133,
2023c.
Xiaoxiao Long, Cheng Lin, Peng Wang, Taku Komura, and Wenping Wang. Sparseneus: Fast gen-
eralizable neural surface reconstruction from sparse views. In European Conference on Computer
Vision, pp. 210–227. Springer, 2022.
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization.
arXiv preprint
arXiv:1711.05101, 2017.
Andrew Luo, Tianqin Li, Wen-Hao Zhang, and Tai Sing Lee. Surfgen: Adversarial 3d shape syn-
thesis with explicit surface discriminators. In Proceedings of the IEEE/CVF International Con-
ference on Computer Vision, pp. 16238–16248, 2021.
Shitong Luo and Wei Hu. Diffusion probabilistic models for 3d point cloud generation. In Proceed-
ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2837–2845,
2021.
Tiange Luo, Chris Rockwell, Honglak Lee, and Justin Johnson. Scalable 3d captioning with pre-
trained models. arXiv preprint arXiv:2306.07279, 2023.
Luke Melas-Kyriazi, Christian Rupprecht, Iro Laina, and Andrea Vedaldi. Realfusion: 360 recon-
struction of any object from a single image. In CVPR, 2023. URL https://arxiv.org/
abs/2302.10663.
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon.
SDEdit: Guided image synthesis and editing with stochastic differential equations. In Interna-
tional Conference on Learning Representations, 2022.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia,
Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision
training. In International Conference on Learning Representations, 2018.
Ben Mildenhall, Pratul P. Srinivasan, Matthew Tancik, Jonathan T. Barron, Ravi Ramamoorthi, and
Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. In ECCV, 2020.
Swaroop Mishra, Daniel Khashabi, Chitta Baral, and Hannaneh Hajishirzi. Cross-task generalization
via natural language crowdsourcing instructions. In Proceedings of the 60th Annual Meeting of
the Association for Computational Linguistics (Volume 1: Long Papers), pp. 3470–3487, 2022.
12
Under review as a conference paper at ICLR 2024
Nasir Mohammad Khalid, Tianhao Xie, Eugene Belilovsky, and Tiberiu Popa. Clip-mesh: Gener-
ating textured meshes from text using pretrained image-text models. In SIGGRAPH Asia 2022
conference papers, pp. 1–8, 2022.
Thomas M¨
uller, Alex Evans, Christoph Schied, and Alexander Keller. Instant neural graphics prim-
itives with a multiresolution hash encoding.
ACM Trans. Graph., 41(4):102:1–102:15, July
2022.
doi: 10.1145/3528223.3530127.
URL https://doi.org/10.1145/3528223.
3530127.
Alex Nichol, Heewoo Jun, Prafulla Dhariwal, Pamela Mishkin, and Mark Chen. Point-e: A system
for generating 3d point clouds from complex prompts. arXiv preprint arXiv:2212.08751, 2022.
Evangelos Ntavelis, Aliaksandr Siarohin, Kyle Olszewski, Chaoyang Wang, Luc Van Gool, and
Sergey Tulyakov. Autodecoding latent 3d diffusion models. arXiv preprint arXiv:2307.05445,
2023.
Roy Or-El, Xuan Luo, Mengyi Shan, Eli Shechtman, Jeong Joon Park, and Ira Kemelmacher-
Shlizerman. StyleSDF: High-Resolution 3D-Consistent Image and Geometry Generation. In
Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),
pp. 13503–13513, June 2022.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong
Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow
instructions with human feedback.
Advances in Neural Information Processing Systems, 35:
27730–27744, 2022.
Dario Pavllo, Graham Spinks, Thomas Hofmann, Marie-Francine Moens, and Aurelien Lucchi.
Convolutional generation of textured 3d meshes. Advances in Neural Information Processing
Systems, 33:870–882, 2020.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten-
hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and
E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research,
12:2825–2830, 2011.
William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint
arXiv:2212.09748, 2022.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas M¨
uller, Joe
Penna, and Robin Rombach. Sdxl: improving latent diffusion models for high-resolution image
synthesis. arXiv preprint arXiv:2307.01952, 2023.
Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d
diffusion. arXiv, 2022.
Guocheng Qian, Jinjie Mai, Abdullah Hamdi, Jian Ren, Aliaksandr Siarohin, Bing Li, Hsin-
Ying Lee, Ivan Skorokhodov, Peter Wonka, Sergey Tulyakov, et al.
Magic123: One image
to high-quality 3d object generation using both 2d and 3d diffusion priors.
arXiv preprint
arXiv:2306.17843, 2023.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal,
Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual
models from natural language supervision. In International conference on machine learning, pp.
8748–8763. PMLR, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi
Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text
transformer. The Journal of Machine Learning Research, 21(1):5485–5551, 2020.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-
conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 1(2):3, 2022.
13
Under review as a conference paper at ICLR 2024
Jeremy Reizenstein, Roman Shapovalov, Philipp Henzler, Luca Sbordone, Patrick Labatut, and
David Novotny. Common objects in 3d: Large-scale learning and evaluation of real-life 3d cat-
egory reconstruction. In Proceedings of the IEEE/CVF International Conference on Computer
Vision, pp. 10901–10911, 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bj¨
orn Ommer. High-
resolution image synthesis with latent diffusion models, 2021.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar
Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic
text-to-image diffusion models with deep language understanding. Advances in Neural Informa-
tion Processing Systems, 35:36479–36494, 2022.
Johannes Lutz Sch¨
onberger, Enliang Zheng, Marc Pollefeys, and Jan-Michael Frahm. Pixelwise
view selection for unstructured multi-view stereo. In European Conference on Computer Vision
(ECCV), 2016.
Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi
Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al. Laion-5b: An
open large-scale dataset for training next generation image-text models.
Advances in Neural
Information Processing Systems, 35:25278–25294, 2022.
Ivan Skorokhodov, Sergey Tulyakov, Yiqun Wang, and Peter Wonka. Epigraf: Rethinking training
of 3d gans. Advances in Neural Information Processing Systems, 35:24487–24501, 2022.
Edward J Smith and David Meger.
Improved adversarial systems for 3d object generation and
reconstruction. In Conference on Robot Learning, pp. 87–96. PMLR, 2017.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna-
tional Conference on Learning Representations, 2021. URL https://openreview.net/
forum?id=St1giarCHLP.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timoth´
ee
Lacroix, Baptiste Rozi`
ere, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and
efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
Alex Trevithick and Bo Yang. Grf: Learning a general radiance field for 3d representation and
rendering. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp.
15182–15192, 2021.
Haochen Wang, Xiaodan Du, Jiahao Li, Raymond A Yeh, and Greg Shakhnarovich. Score jaco-
bian chaining: Lifting pretrained 2d diffusion models for 3d generation. In Proceedings of the
IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12619–12629, 2023a.
Qianqian Wang, Zhicheng Wang, Kyle Genova, Pratul Srinivasan, Howard Zhou, Jonathan T. Bar-
ron, Ricardo Martin-Brualla, Noah Snavely, and Thomas Funkhouser. Ibrnet: Learning multi-
view image-based rendering. In CVPR, 2021.
Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu. Prolific-
dreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. arXiv
preprint arXiv:2305.16213, 2023b.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, An-
drew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. In International
Conference on Learning Representations, 2021.
Jiajun Wu, Chengkai Zhang, Tianfan Xue, Bill Freeman, and Josh Tenenbaum. Learning a proba-
bilistic latent space of object shapes via 3d generative-adversarial modeling. Advances in neural
information processing systems, 29, 2016.
Alex Yu, Vickie Ye, Matthew Tancik, and Angjoo Kanazawa. pixelNeRF: Neural radiance fields
from one or few images. In CVPR, 2021.
14
Under review as a conference paper at ICLR 2024
Xiaohui Zeng, Arash Vahdat, Francis Williams, Zan Gojcic, Or Litany, Sanja Fidler, and Karsten
Kreis. Lion: Latent point diffusion models for 3d shape generation. In Advances in Neural
Information Processing Systems (NeurIPS), 2022.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable
effectiveness of deep features as a perceptual metric. In CVPR, 2018.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat,
Ping Yu, Lili Yu, et al. Lima: Less is more for alignment. arXiv preprint arXiv:2305.11206, 2023.
Linqi Zhou, Yilun Du, and Jiajun Wu. 3d shape generation and completion through point-voxel
diffusion. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp.
5826–5835, 2021.
A
APPENDIX
A.1
DIVERSITY OF GENERATION.
Our method can generate diverse results from the same text prompt by using different random seeds
in the feed-forward pass, where optimization-based methods are known As shown in Figure 6, our
approach excels in generating diverse 3D assets featuring strikingly distinct textures and geometries
from the same prompt.
A.2
MULTI-VIEW DIFFUSION TRAINING DETAILS
We use SDXL as the base model for our first-stage fine-tuning. We use the AdamW optimizer with a
fixed learning rate 10−5, β1 = 0.9, β2 = 0.999 and a weight decay of 10−2. We fine-tune the model
using fp16 on 32 A100 GPUs with a total batch size of 192. No gradient accumulation is used. We
train the model on 10K curated data for 40K steps, which takes around 3 hours.
We train the model with the standard denoising diffusion loss Ho et al. (2020)
L(θ) = Et,x0,ϵ
∥ϵ −ϵθ(√αtx0 +
√
1 −αtϵ, t)∥2
(1)
where ϵθ is the denoising UNet and θ are the trainable parameters.
The SDXL introduces image resolution and aspect ratio conditioning that allow mixing training on
images of different resolutions and aspect ratios. As for our training data, we render 4 views each
with a resolution of 512x512 and assemble them into a 1024x1024 images. Therefore we fix the
resolution and aspect ratio conditioning to be (1024, 1024) throughout the fine-tuning procedure.
We don’t do random cropping in our training and fixed the crop conditioning to be (0, 0). All the
other training setups are identical to the SDXL training.
A.3
DATA CURATION DETIALS
We adopt a quality scorer to curate high-quality data from the Objaverse dataset. To train the quality
scorer, we first randomly sample 2000 data points from the dataset and manually label each 3D asset
as good or bad. Good assets have realistic textures and complex geometry, while bad ones have
simple shapes and flat or cartoon-like texture. This criterion is subjective and imprecise, but we
found it to be good enough for the purpose of data filtering.
Since the amount of annotated data is limited, we use a pretrained CLIP (Radford et al., 2021) model
to extract high-level image features of rendered images at 5 randomly sampled camera viewpoints
for each asset. Then we train a simple binary SVM classifier on top of the averaged CLIP features
from different views. We use the NuSVC implementation from the popular sklearn framework
Pedregosa et al. (2011), which also gives us probability estimates of the classification.
After training the SVM model, we compute the scores for all the data points in the dataset. Similar
to the training data, for each 3D object we render 5 random views, compute the averaged CLIP
features, and use the trained SVM to estimate the probability values. These probabilities are used
as scores to rank the data from high to low quality. Finally, we use the top 10K objects as our
fine-tuning data.
15
Under review as a conference paper at ICLR 2024
Figure 6: Our method can generate diverse results from the same text prompt.
To render the 4-view data, we scale the curated objects and center them at a cube [−1, 1]3. We
render the objects with a white background following the structured setup discussed in Section 3.1
using a field view of 50◦at a distance of 2.7 under uniform lighting. We use the physically-based
path tracer Cycles in Blender for rendering.
In Figure 7 we show qualitative comparisons on results from models trained with curated data and
random data. Models trained with random data tend to generate cartoon-like 3D assets with simple
and flat textures. This is not surprising since a bulk of the Objaverse dataset contains simple shapes
with simple textures, and without curation these data will guide the model to over denoise the results,
leading to large areas of flat color. On the contrary, models trained with curated data tend to generate
realistically looking objects with complex textures and shading effects.
A.4
SPARSE-VIEW RECONSTRUCTION DETAILS
Model Details
We use the DINO-ViT-B/16 as our image encoder. This model is transformer-
based, which has 12 layers and the hidden dimension of the transformer is 768. The ViT begins
with a convolution of kernel-size 16, stride 16, and padding 0. It is essentially patchifying the input
image with a patch size of 16 × 16. For our final model, the image input resolution is 512, thus it
leads to 32 × 32 = 1024 spatial tokens in the vision transformer. In ablation studies, we reduce the
input resolution from 512 to 256 to save some compute budget. The orignal DINO is trained with
a resolution of 224, thus the positional embedding has only a size of 14 × 14 = 196. We thus use
2D bilinear extrapolation (with the torch.nn.functional.interpolate function) to extrapolate it to the
desired token size.
Since the camera information is integrated into the image encoder via modulation, we inject mod-
ulation layer (Peebles & Xie, 2022) into each of the transformer layer (for both self-attention layer
and the MLP layer). The modulation layer is initialized to be an identity mapping thus it is suitable
to be added to a pre-trained vision transformer.
After the image encoder, we have 1025 image feature tokens for each image, since we also include
the output of the [CLS] token. We concatenate the tokens from all four images to construct a
sequence condition features of length 4100. This condition feature will be used to create the keys
and values in the cross-attention layer of the image-to-triplane transformer decoder.
The image-to-triplane transformer decoder starts with a token sequence of (3 × 32 × 32) × 1024,
where (3 × 32 × 32) is the number of tokens and 1024 is the hidden dimension of the transformer.
We use 16 layers in our transformer decoder. All attention layer has 16 attention heads and each
head has a dimension of 64. We remove the bias term in the attention layer as in Touvron et al.
(2023). We take the pre-normalization architecture of the transformer where each sub-layer will be
in the format of x + f(LayerNorm(x)).
After the transformer, we apply the de-convolution layer to map the transformer output from (3 ×
32 × 32) × 1024 to 3 × (64 × 64) × 80. It means that there are 3 planes (XY, YZ, XZ) Chan et al.
(2022) and each plane has a size of 64 by 64. The dimension of each plane is 80 All three planes
share the same deconvlution layer. The deconvolution is of kernal size 2, stride 2, and pad 0.
In NeRF volumetric rendering, the features from the three planes are bilinearly interpolated and also
concatenated to get a 240 dimensional feature for each point. Then, we have 10 layers MLP with a
hidden dimension of 64 to map this 240-dim feature to a 4-dim feature. The first three dimension
will be treated as RGB and normalized to (0, 1) with a sigmoid function. The last dimension will
16
Under review as a conference paper at ICLR 2024
be treated as the density value of the NeRF field. Since density is non-negative, we use exponential
function to map the MLP’s output to be non-negative.
For the exact formulation of the above operators, please refer to LRM (Anonymous, 2023) and
DiT (Peebles & Xie, 2022).
Training Details
We use the AdamW (Kingma & Ba, 2014; Loshchilov & Hutter, 2017) optimizer
to train our model. We use a peak learning rate of 4 × 10−4 with a linear warm-up (on the first 3K
steps) and a cosine decay. We change the β2 of the AdamW optimizer to 0.95 for better stability. We
use a weight-decay of 0.05 for non-bias and non-layernorm parameters. We use a gradient clipping
of 1.
The training is completed with 120 epochs on our rendered Objaverse data with a training batch size
of 1024. We use both L2 loss and LPIPS loss to supervise the model and the weight of two losses
are L2=1 and LPIPS=2. The model is trained on 128 NVIDIA A100 GPUs and the whole training
can be finished in 7 days.
For the initialization of the image encoder, we initialize it from the official-released DINO pre-
training weight. For the initialization of the triplane decoder, and NeRF MLP, we use the default
initializer in the PyTorch implementation. We empirically found that the pre-normalization trans-
former is robust to different initialization of linear layers. For the positional embedding of the
triplane tokens in the transformer decoder, we initialize with a Gaussian random of zero-mean and
std of 1/
√
1024.
For each training step, we take 4 random views for input and 4 random views as supervision. To
allow the input views and supervision views to have more overlapping than purely independent
random samples, we first sample 8 random views from all views. Then the 4 input views and 4
supervising views are randomly sampled from these 8 views.
The number of points per ray in NeRF rendering is 128, which is uniformly distributed inside the
[−1, 1]3 object bounding box. The rendering resolution is 128 × 128. To allow higher actual super-
vising resolution, we first resize the image to a smaller resolution (uniformly sampled from [128,
384]) and then crop a patch of 128 × 128 from it. Thus we can go beyond the rendering resolution
of 128.
We utilize the flash attention (Dao et al., 2022), mixed-precision training (with bf16 as the half-
precision format) (Micikevicius et al., 2018), and gradient checkpointing (Chen et al., 2016) to
improve the compute/memory efficiency of the training.
A.5
STABLE DIFFUSION 1.5 (SD1.5) TRAINING DETAILS
We use 8 A100 GPUs for fine-tuning SD1.5 on 100K data with a total batch size of 64. We use the
same AdamW optimizer as on SDXL with exactly the same hyper-parameters. We also use gradient
accumulation for 3 steps, which gives an effective batch size of 192. The training loss is the same as
SDXL. We train the model for 120K steps (40K parameter updates due to gradient accumulation),
which takes roughly 33 hours.
A.6
EXTENSION TO IMAGE-CONDITIONED GENERATION
Our method can also be extended to support additional image conditioning to provide more fine-
grained control over the 3D model to be generated. In this process, the input to the model includes an
input text prompt that describes the object to be generated as well as a single image of the object. We
use the same training data as our text-conditioned model. During training, for a randomly sampled
time step, we aim to keep the first image (top-left quadrant) untouched and only add noise to the
remaining three views. This allows the diffusion model to generate the other views while accounting
for the conditioning image. Figure 8 shows some visual results of our image-conditioned model.
From the results we can see that our method is able to effectively generate the other views with
faithful details that are coherent with the input text prompt and image, thus giving us a high-quality
3D model.
During training, suppose we have a clean 2 × 2 grid image with four views of the same object.
In normal diffusion training, we add noise to the entire image and feed it to the UNet. To allow
17
Under review as a conference paper at ICLR 2024
Table 4: Ablation study of the sparse-view reconstruction model.
#Layers
Render
Supervision
PSNR ↑
SSIM ↑
LPIPS ↓
exp01
6
64
All
23.6551
0.8616
0.1281
exp02
12
64
All
23.8257
0.8631
0.1266
exp03
24
64
All
23.8351
0.8635
0.1258
exp04
12
32
All
23.1704
0.8561
0.1358
exp05
12
64
w/o novel
18.2359
0.8103
0.2256
exp06
12
64
w/o LPIPS
24.1699
0.8641
0.1934
the model to condition on a single image (the first view), we keep that left upper part of the image
always clean, and only add noise the other parts of the grid image. This kind of input is different from
normal inputs to the original diffusion model, but after finetuning the model can learn to adapt to
the distribution shift and be able to denoise conditioned on the upper left image. During inference,
similarly, we replace the upper left part of every intermediate iterate with the clean conditioned
image.
A.7
GAUSSIAN BLOBS INITIALIZATION
Since the diffusion model is finetuned with only a relatively small number of steps, it still largely
possesses the original denoising behavior on images that are not in the form of 2×2 grids and do not
have white background. Naively applying the standard backward denoising process starting from
random Gaussian noise will likely lead to results far from the data distribution of the finetuning data
(see Figure 5).
The spatial structure of the training images is simple: four views of the same object are placed at
the centers of each quadrant. Also, the background is always white. Since the model is finetuned on
such data with a denoising objective, it is naturally that, when presented with a noisy input whose
underlying clean image has these two characteristics, the model will tend to denoise the image to
a clean one where the four quadrant objects are view consistent. Following this, and inspired by
SDEdit Meng et al. (2022), we introduce Gaussian blobs initialization to bias the model towards
generating samples consistent with the distribution of the fine-tuning data.
The standard latent diffusion inference starts with a Gaussian noise image ϵ with the same size
as the image latents. Instead, we modify the initial iterate to be a composition of Gaussian noise
and an image with the aforementioned two characteristics: object quadrants and white background.
We construct such an image by first constructing a white image with a black Gaussian blob at the
center. Specifically, we construct a H × W grayscale image I, where H and W are the height and
width of the input RGB image with value range [0, 1]. For all our models H = W, and we denote
S = H = W. For a given pixel (x, y), the value of I is computed as
I(x, y) = 1 −exp
−(x −S/2)2 + (y −S/2)2
2σ2S2
(2)
where σ is a hyper-parameter controlling the width of the gaussian blob. Such an image looks like
a black ball at the center of a white image slowly fading away towards the edges of the image. We
then assemble four such same images into a 2 × 2 image grid. Some examples of such images with
different σ can be seen at the first row of figure 5.
Next we construct the initial noise for the denoising step by blending a complete Gaussian noise
latent with the latent of the Gaussian blobs latent. We denote the latents of I as ˜
I, and a noise latent
image with i.i.d. Gaussian values as ϵ. For a N step denoising inference process with timesteps
{tN, tN−1, ..., t0}, we mix the two latents with a weighted sum
ϵtN =
p
αtN ˜
I +
p
1 −αtN ϵ
(3)
Then ϵtN is used as the initial noise of the denoising process. E.g., the tN is 980 for a denoising
step with 50 (and the total time step is 1000).
A.8
SPARSE VIEW RECONSTRUCTION ABLATION RESULTS
We conduct the ablation study of our second-stage model (the sparse-view reconstruction model)
on the same training set as our final model. However, we change the training recipe to reduce the
18
Under review as a conference paper at ICLR 2024
computation cost to 32 A100 GPU for 1 day. The change of configuration for ablation is: (1) input
image resolution of 256 (2) 96 points per ray during rendering (3) 5 layers instead of 10 layers in
NeRF MLP. (4) number of epochs decreased to 30.
To evaluate the performance of different variants, we test them on another 3D dataset Google
Scanned Object (GSO) (Downs et al., 2022). For each object in GSO, we render a set of 64 view
images with resolution 512 × 512 and white background from elevations 0◦, 20◦, 40◦, 60◦. Each
elevation has 16 views with equidistant azimuths starting from 0. We use 4 views on elevation 20◦
degrees and azimuths 45◦, 135◦, 225◦, 315◦as input views, and randomly sample 5 views from the
other 60 views as ground truth images. We fix the random seed across different runs to make sure
that the selected GT views are the same. A model takes in 4 input images and render 5 512 × 512
images according to the camera parameters of the 5 GT views. Then we compute 3 metrics between
the rendered images and the corresponding GT images to measure the performance including PSNR,
SSIM and LPIPS. These metrics are averaged over all of the 1019 objects in the dataset to compute
the final values.
The results of the ablation studies are in Figure 4. First, we show that the model is robust to the
transformer hyperparameters (i.e., the transformer layers) in exp01, exp02, and exp03. Second, we
show that the LPIPS loss can largely change the results by comparing the exp02 and exp06. The
only difference between these two experiments is the use of LPIPS loss during training. Without
LPIPS loss, the model drops a lot on LPIPS while got slightly improvement on PSNR and SSIM.
However, we empirically find that LPIPS is much more aligned with human judgement and the
rendered images become blurry without LPIPS. The rendering resolution is also important (as shown
in exp04) because of the LPIPS loss, since LPIPS can be more robust and accurate for larger image
resolution. Thus we use a rendering resolution of 128 by 128 in training our final model.
Also, the inclusion of novel supervising views in the training is also important as shown in exp05. All
three metrics got a significant drop when only supervising on the four input views. After checking
the results, the reason is that the coverage of four views is usually not complete for the whole space.
Thus the model will get grey floaters at the space that are not cover by these four views.
B
LIMITATIONS AND DISCUSSION
While our model outperforms existing works in terms of speed, its efficiency is ultimately con-
strained by the underlying 2D diffusion models. We have also empirically observed that the fine-
tuned multi-view generation model necessitates more inference steps compared to its original 2D
diffusion counterpart. Another limitation is the absence of a 3D inductive bias in the initial stage
of the model; consequently, it is wholly dependent on a 2D model to generate 3D-consistent views,
which could be a limiting factor. On the other hand, the feed-forward multi-view reconstruction
model tends to blur textures, resulting in a slight degradation of the reconstructed 3D model’s qual-
ity in comparison to the original views.
19
Under review as a conference paper at ICLR 2024
Curated
Random
a hippo wearing a
sweater
a panda wearing a
necktie and sitting
in an office chair
a stack of pancakes
covered in maple syrup
a gummy bear
driving a convertible
Figure 7: Comparisons on novel view renderings of NeRF assets generated from SDXL models
fine-tuned with 10K curated and random data. We can see that that curated data enables the model
to generate more photorealistic 3D assets with more geometric and textures details. Here curated
and random correspond to Exp d (Curated-10K-s10K) and i (Random-10K-s10K) in Table 3.
20
Under review as a conference paper at ICLR 2024
Input image
Our 2x2 grid
Ours
Zero-1-to-3
One-2-3-45
Figure 8: Comparison to previous methods on single image-conditioned 3D generation. We com-
pared to previous methods Zero-1-to-3 (Liu et al., 2023b) and One-2-3-45 (Liu et al., 2023a). Our
method can faithfully generate the details at the invisible regions, thus empowering us to reconstruct
3D assets of higher quality than baseline methods. All input images are generated with a public
text-to-image platform Adobe Firefly (Adobe, 2023).
21
Under review as a conference paper at ICLR 2024
ppendix blobs
Sigma=0.0
Sigma=0.01
Sigma=0.1
Sigma=0.2
Sigma=0.3
Figure 9: 2x2 grid images generated with Gaussian blobs of different Sigma σ.
22
Under review as a conference paper at ICLR 2024
10 steps
(6 seconds)
20 steps
(8 seconds)
50 steps
(15 seconds)
100 steps
(20 seconds)
a blue poison-dart frog sitting on a water lily
beautifully carved jade heavenly Chinese guardian lion, hand carved nephrite jade
a cute dragon with big eyes
Figure 10: Comparison on the NeRF assets generated with different number of DDIM steps and
their inference time. While we use 100 steps in our experiments that takes 20 seconds to generate
a NeRF asset, we find that using smaller number of steps can also give us results of similar quality
with much shorter inference time.
23
Under review as a conference paper at ICLR 2024
a plush toy of a corgi nurse
a chimpanzee holding a cup of hot coffee
a hippo made out of chocolate
a red-eyed tree frog
Ours
ProlificDreamer
Figure 11: SDS optimization-based method such as ProlificDreamer (Wang et al., 2023b) can possi-
bly suffer from the Janus problem, which greatly degrades the quality of the 3D assets. In contrast,
our method can effectively get rid of this problem.
a tray of sushi
containing pugs
a squirrel dressed
like a clown
baby elephant jumping
on a trampoline
a bear dressed as a
lumberjack
SDXL
SD1.5
Figure 12: Comparisons on the quality of the NeRF assets generated with fine-tuned SDXL and
SD1.5 models. SDXL has a model size that is three times larger than SD1.5 thus has better text
comprehension. As shown in the figure, the 3D assets generated by our fine-tuned SDXL has better
photo realism and text alignment. The used SDXL and SD1.5 models are from Exp d (Curated-10K-
s10K) and m (Curated-100K-s40K) in Table 3.
24
Under review as a conference paper at ICLR 2024
1K
10K
100K
a brightly colored
mushroom
growing on a log
a dalmatian wearing
a fireman’s hat
a teal moped
a pug wearing a bee
costume
Figure 13: Comparison on the effect of different fine-tuning data size. Training on too little data such
as 1K results in inconsistency between the generated 4 views, thus resulting in incorrect geometry.
On the other side, training on too much data such as 100K makes the model bias towards the fine-
tuning dataset, thus negatively affecting the quality of generated 3D assets. Here 1K, 10K and 100K
are corresponding to Exp a (Curated-1K-s1K), d (Curated-10K-s10K) and g (Curated-100K-s40K)
in Table 3, respectively.
a beagle eating a
donut
a bear dancing
ballet
a blue motorcycle
a chihuahua
wearing a tutu
20K steps
10K steps
4K steps
Figure 14: Comparison on different number of fine-tuning steps. 4K training steps lead to inconsis-
tent 4-view generation, while 20K results in biasing towards the fine-tuning data. In contrast, 10K
achieves a balance between these two. Here 4K, 10K and 20K correspond to Exp c (Curated-10K-
s4K), d(Curated-10K-s10K) and e (Curated-10K-s20K) in Table 3.
25choice A
In the repo, OpenLRM’s reliance on large datasets like Objaverse and MVImgNet introduces challenges in balancing reconstruction fidelity with computational efficiency, particularly when scaling to more complex scenes compared to the more efficient architecture of Instant3D, which is optimized for real-time applications and lower-latency tasks, but this leads to a tradeoff in latency and scalability.
choice B
From the python files, we can see that while Instant3D focuses heavily on computational efficiency for rapid object generation in real-time, OpenLRM’s large-scale reconstruction model prioritizes generalization across various 3D environments. However, the hybrid datasets, Objaverse and MVImgNet, introduce issues with model overfitting to certain object types, leading to difficulties in scaling and maintaining real-time generation accuracy, unlike Instant3D’s streamlined approach.
choice C
In the core openlrm package, we see that OpenLRM’s approach focuses on generating high-fidelity 3D reconstructions using hybrid datasets, which can be computationally demanding. However, unlike Instant3D’s optimization for speed and lower resource usage, OpenLRM faces challenges in optimizing its memory consumption and real-time performance when scaling to larger datasets or more detailed reconstructions.
choice D
We can note from the codebase that the architectural difference between OpenLRM and Instant3D lies in their treatment of real-time constraints and scalability. OpenLRM is designed to reconstruct larger scenes and complex objects, which Instant3D resolves by focusing more narrowly on specific object categories and reducing the training dataset size to improve real-time performance.
difficulty
hard
domain
Code Repository Understanding
length
medium
sub domain
Code repo QA
Discussion
No discussion posts on this page yet. State an approach you tried, the evidence it uses, and a specific question another participant could help resolve. Use the posting template.
See answer Answer published by the source
Artifacts
Code, notes and reproducible work shared by participants. Files are served from a separate origin.
No artifacts on this page yet. Share reproducible code or notes in a contribution. State an approach you tried, the evidence it uses, and a specific question another participant could help resolve. Use the posting template.
Source and history
initial import