example.py

import os
import sys

import argparse
import time
import imageio
import numpy as np
from PIL import Image

sys.path.insert(0, './hy3dshape')
os.environ['TORCH_CUDA_ARCH_LIST'] = '9.0'
os.environ['ATTN_BACKEND'] = 'xformers'
os.environ['SPCONV_ALGO'] = 'native'

from trellis.pipelines import NeARImageToRelightable3DPipeline
from hy3dshape.pipelines import Hunyuan3DDiTFlowMatchingPipeline


def main():
    parser = argparse.ArgumentParser(description='NeAR: from image or SLaT to relightable 3D')
    parser.add_argument('--checkpoint', type=str, default='checkpoints',
                        help='Pipeline weights directory (contains pipeline.yaml)')
    parser.add_argument('--image', type=str, default="assets/example_image/T.png",
                        help='Input image path (one of --slat or --image)')
    parser.add_argument('--slat', type=str, default=None,
                        help='SLaT .npz path (one of --image or --slat)')
    parser.add_argument('--hdri', type=str, default="assets/hdris/studio_small_03_1k.exr",
                        help='HDRI .exr path')
    parser.add_argument('--out_dir', type=str, default='relight_out',
                        help='Output directory')
    parser.add_argument('--yaw', type=float, default=0.0, help='View yaw (degrees) ')
    parser.add_argument('--pitch', type=float, default=0.0, help='View pitch (degrees)')
    parser.add_argument('--hdri_rot', type=float, default=0.0,
                        help='HDRI rotation angle (degrees)')
    parser.add_argument('--video_frames', type=int, default=40,
                        help='Render additional spiral camera path video frames')
    parser.add_argument('--seed', type=int, default=42, help='Random seed')
    parser.add_argument('--save_slat', type=str, default=None,
                        help='When generating from image, save SLaT to this .npz path')
    parser.add_argument('--no_cuda', action='store_true', help='Do not use CUDA')
    args = parser.parse_args()

    if (args.image is None) == (args.slat is None):
        parser.error('Please specify --image or --slat, one of them')
    from_image = args.image is not None

    os.makedirs(args.out_dir, exist_ok=True)
    device = 'cuda' if not args.no_cuda else 'cpu'

    total_t0 = time.perf_counter()
    t0 = time.perf_counter()
    hyshape_model_id = 'tencent/Hunyuan3D-2.1'
    hyshape_pipe = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(hyshape_model_id)
    hyshape_pipe.to(device)
    pipeline = NeARImageToRelightable3DPipeline.from_pretrained(args.checkpoint)
    pipeline.to(device)
    print(f"  [OK] Pipeline loaded, +{time.perf_counter() - t0:.1f}s")

    if from_image:
        image = Image.open(args.image).convert('RGB')
        image_prep = pipeline.preprocess_image(image)

        t0 = time.perf_counter()
        mesh = hyshape_pipe(image=image_prep)[0]
        mesh_path = os.path.join(args.out_dir, 'initial_3d_shape.glb')
        mesh.export(mesh_path)
        print(f"  [OK] Geometry mesh generated, +{time.perf_counter() - t0:.1f}s, saved to {mesh_path}")

        t0 = time.perf_counter()
        slat = pipeline.run_with_shape(
            image_prep,
            mesh,
            seed=args.seed,
            preprocess_image=False,
        )
        print(f"  [OK] Image → SLaT generated, +{time.perf_counter() - t0:.1f}s")
        if args.save_slat:
            np.savez(
                args.save_slat,
                feats=slat.feats.cpu().numpy(),
                coords=slat.coords.cpu().numpy(),
            )
            print(f"  [OK] Saved SLaT to: {args.save_slat}")
    else:
        t0 = time.perf_counter()
        slat = pipeline.load_slat(args.slat)
        print(f"  [OK] Loaded SLaT, +{time.perf_counter() - t0:.1f}s")

    t0 = time.perf_counter()
    hdri_np = pipeline.load_hdri(args.hdri)
    print(f"  [OK] Loaded HDRI, +{time.perf_counter() - t0:.1f}s")

    t0 = time.perf_counter()

    pipeline.renderer.ssaa = 1
    pipeline.renderer.resolution = 1024

    views = pipeline.render_view(
        slat, hdri_np,
        yaw_deg=args.yaw, pitch_deg=args.pitch,
        fov=40.0, radius=2.0, hdri_rot_deg=args.hdri_rot,
        resolution=512
    )

    color_path = os.path.join(args.out_dir, 'relight_color.png')
    base_color_path = os.path.join(args.out_dir, 'base_color.png')
    metallic_path = os.path.join(args.out_dir, 'metallic.png')
    roughness_path = os.path.join(args.out_dir, 'roughness.png')
    shadow_path = os.path.join(args.out_dir, 'shadow.png')
    views['color'].save(color_path)
    views['base_color'].save(base_color_path)
    views['metallic'].save(metallic_path)
    views['roughness'].save(roughness_path)
    views['shadow'].save(shadow_path)

    print(f"  [OK] Single view rendering completed, +{time.perf_counter() - t0:.1f}s, saved to {color_path}, {base_color_path}, {metallic_path}, {roughness_path}, {shadow_path}")

    if args.video_frames > 0:
        t0 = time.perf_counter()
        frames = pipeline.render_camera_path_video(
            slat, hdri_np, num_views=args.video_frames,
            fov=40.0, radius=2.0, hdri_rot_deg=args.hdri_rot,
            verbose=True, full_video=True, shadow_video=True
        )
        video_path = os.path.join(args.out_dir, 'relight_camera_path.mp4')
        imageio.mimsave(video_path, frames, fps=24)
        print(f"  [OK] Camera path video completed, +{time.perf_counter() - t0:.1f}s, saved to {video_path}")
    print("Done.")

    if args.video_frames > 0:
        t0 = time.perf_counter()
        hdri_roll_frames, render_frames = pipeline.render_hdri_rotation_video(
            slat, hdri_np, num_frames=args.video_frames,
            yaw_deg=args.yaw, pitch_deg=args.pitch,
            fov=40.0, radius=2.0,
            verbose=True, full_video=True, shadow_video=True
        )
        hdri_video_path = os.path.join(args.out_dir, 'hdri_roll.mp4')
        render_video_path = os.path.join(args.out_dir, 'relight_hdri_rotation.mp4')
        imageio.mimsave(hdri_video_path, hdri_roll_frames, fps=24)
        imageio.mimsave(render_video_path, render_frames, fps=24)
        print(
            f"  [OK] HDRI rotation videos completed, +{time.perf_counter() - t0:.1f}s, "
            f"saved to {hdri_video_path} and {render_video_path}"
        )

    print(f"  [OK] Total time: {time.perf_counter() - total_t0:.1f}s")


if __name__ == '__main__':
    main()