This repository contains the official code for the paper Neural LightRig: Unlocking Accurate Object Normal and Material Estimation with Multi-Light Diffusion.

git clone https://github.com/ZexinHe/Neural-LightRig.git
cd Neural-LightRig

We recommend python>=3.11 and torch>=2.4.0. Not all versions are tested, but the code should work with other versions as well.

Please make sure torch is ready before installing the following requirements.

pip install -r requirements.txt

Pretrained models are available on Hugging Face. They will be automatically downloaded when running inference for the first time.

Images should be background-removed in advance and put into a folder. We provide example RGBA images in ./assets/examples.

Run the following command to perform inference on the example images. The results will be saved in ./results/example_outputs. Please modify the --img_dir and --save_dir arguments to your own paths.

python inference.py --img_dir "./assets/examples" --save_dir "./results/example_outputs"

More inference arguments are available, including seed, classifer-free-guidance, inference steps, and more. Please refer to inference.py for more details.

Our rendering dataset LightProp is on Hugging Face. Please follow instructions there on how to download and prepare the dataset.

Update training configs. Run training with the following example command.

torchrun --nproc_per_node=1 -m mld.train --base mld/configs/train-phase-1.yaml

Then update checkpoint path for the trained model in phase-2 config. Continue training with the following example command.

torchrun --nproc_per_node=1 -m mld.train --base mld/configs/train-phase-2.yaml

Update training configs for 256-resolution. You may either train from scratch or finetune on our checkpoint for your own purposes. Run the following example command to start training.

torchrun --nproc_per_node=1 -m recon.trainer --group NLR --case train-256

Update checkpoint path for the trained 256-resolution model in 512-resolution config. Continue training with the following example command.

torchrun --nproc_per_node=1 -m recon.trainer --group NLR --case train-512

We thank the authors of the following repositories for their great works!

• diffusers
• InstantMesh

If you find this work useful for your research, please consider citing

@misc{neural_lightrig,
    title={Neural LightRig: Unlocking Accurate Object Normal and Material Estimation with Multi-Light Diffusion}, 
    author={Zexin He and Tengfei Wang and Xin Huang and Xingang Pan and Ziwei Liu},
    year={2024},
    eprint={2412.09593},
    archivePrefix={arXiv},
    primaryClass={cs.CV},
    url={https://arxiv.org/abs/2412.09593},
}