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()