LTT/PRM
24
1import os2import argparse3import glm4import numpy as np5import torch6import rembg7from PIL import Image8from torchvision.transforms import v29import torchvision10from pytorch_lightning import seed_everything11from omegaconf import OmegaConf12from einops import rearrange, repeat13from tqdm import tqdm14from huggingface_hub import hf_hub_download15from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler16 17from src.data.objaverse import load_mipmap18from src.utils import render_utils19from src.utils.train_util import instantiate_from_config20from src.utils.camera_util import (21 FOV_to_intrinsics, 22 center_looking_at_camera_pose,23 get_zero123plus_input_cameras,24 get_circular_camera_poses,25)26from src.utils.mesh_util import save_obj, save_obj_with_mtl27from src.utils.infer_util import remove_background, resize_foreground, save_video28 29def str_to_tuple(arg_str):30 try:31 return eval(arg_str)32 except:33 raise argparse.ArgumentTypeError("Tuple argument must be in the format (x, y)")34 35 36def get_render_cameras(batch_size=1, M=120, radius=4.0, elevation=20.0, is_flexicubes=False, fov=50):37 """38 Get the rendering camera parameters.39 """40 train_res = [512, 512]41 cam_near_far = [0.1, 1000.0]42 fovy = np.deg2rad(fov)43 proj_mtx = render_utils.perspective(fovy, train_res[1] / train_res[0], cam_near_far[0], cam_near_far[1])44 all_mv = []45 all_mvp = []46 all_campos = []47 if isinstance(elevation, tuple):48 elevation_0 = np.deg2rad(elevation[0])49 elevation_1 = np.deg2rad(elevation[1])50 for i in range(M//2):51 azimuth = 2 * np.pi * i / (M // 2)52 z = radius * np.cos(azimuth) * np.sin(elevation_0)53 x = radius * np.sin(azimuth) * np.sin(elevation_0)54 y = radius * np.cos(elevation_0)55 56 eye = glm.vec3(x, y, z)57 at = glm.vec3(0.0, 0.0, 0.0)58 up = glm.vec3(0.0, 1.0, 0.0)59 view_matrix = glm.lookAt(eye, at, up)60 mv = torch.from_numpy(np.array(view_matrix))61 mvp = proj_mtx @ (mv) #w2c62 campos = torch.linalg.inv(mv)[:3, 3]63 all_mv.append(mv[None, ...].cuda())64 all_mvp.append(mvp[None, ...].cuda())65 all_campos.append(campos[None, ...].cuda())66 for i in range(M//2):67 azimuth = 2 * np.pi * i / (M // 2)68 z = radius * np.cos(azimuth) * np.sin(elevation_1)69 x = radius * np.sin(azimuth) * np.sin(elevation_1)70 y = radius * np.cos(elevation_1)71 72 eye = glm.vec3(x, y, z)73 at = glm.vec3(0.0, 0.0, 0.0)74 up = glm.vec3(0.0, 1.0, 0.0)75 view_matrix = glm.lookAt(eye, at, up)76 mv = torch.from_numpy(np.array(view_matrix))77 mvp = proj_mtx @ (mv) #w2c78 campos = torch.linalg.inv(mv)[:3, 3]79 all_mv.append(mv[None, ...].cuda())80 all_mvp.append(mvp[None, ...].cuda())81 all_campos.append(campos[None, ...].cuda())82 else:83 # elevation = 90 - elevation84 for i in range(M):85 azimuth = 2 * np.pi * i / M86 z = radius * np.cos(azimuth) * np.sin(elevation)87 x = radius * np.sin(azimuth) * np.sin(elevation)88 y = radius * np.cos(elevation)89 90 eye = glm.vec3(x, y, z)91 at = glm.vec3(0.0, 0.0, 0.0)92 up = glm.vec3(0.0, 1.0, 0.0)93 view_matrix = glm.lookAt(eye, at, up)94 mv = torch.from_numpy(np.array(view_matrix))95 mvp = proj_mtx @ (mv) #w2c96 campos = torch.linalg.inv(mv)[:3, 3]97 all_mv.append(mv[None, ...].cuda())98 all_mvp.append(mvp[None, ...].cuda())99 all_campos.append(campos[None, ...].cuda())100 all_mv = torch.stack(all_mv, dim=0).unsqueeze(0).squeeze(2)101 all_mvp = torch.stack(all_mvp, dim=0).unsqueeze(0).squeeze(2)102 all_campos = torch.stack(all_campos, dim=0).unsqueeze(0).squeeze(2)103 return all_mv, all_mvp, all_campos104 105def render_frames(model, planes, render_cameras, camera_pos, env, materials, render_size=512, chunk_size=1, is_flexicubes=False):106 """107 Render frames from triplanes.108 """109 frames = []110 albedos = []111 pbr_spec_lights = []112 pbr_diffuse_lights = []113 normals = []114 alphas = []115 for i in tqdm(range(0, render_cameras.shape[1], chunk_size)):116 if is_flexicubes:117 out = model.forward_geometry(118 planes,119 render_cameras[:, i:i+chunk_size],120 camera_pos[:, i:i+chunk_size],121 [[env]*chunk_size],122 [[materials]*chunk_size],123 render_size=render_size,124 )125 frame = out['pbr_img']126 albedo = out['albedo']127 pbr_spec_light = out['pbr_spec_light']128 pbr_diffuse_light = out['pbr_diffuse_light']129 normal = out['normal']130 alpha = out['mask']131 else:132 frame = model.forward_synthesizer(133 planes,134 render_cameras[i],135 render_size=render_size,136 )['images_rgb']137 frames.append(frame)138 albedos.append(albedo)139 pbr_spec_lights.append(pbr_spec_light)140 pbr_diffuse_lights.append(pbr_diffuse_light)141 normals.append(normal)142 alphas.append(alpha)143 144 frames = torch.cat(frames, dim=1)[0] # we suppose batch size is always 1145 alphas = torch.cat(alphas, dim=1)[0] 146 albedos = torch.cat(albedos, dim=1)[0]147 pbr_spec_lights = torch.cat(pbr_spec_lights, dim=1)[0]148 pbr_diffuse_lights = torch.cat(pbr_diffuse_lights, dim=1)[0]149 normals = torch.cat(normals, dim=0).permute(0,3,1,2)[:,:3]150 return frames, albedos, pbr_spec_lights, pbr_diffuse_lights, normals, alphas151 152 153###############################################################################154# Arguments.155###############################################################################156 157parser = argparse.ArgumentParser()158parser.add_argument('config', type=str, help='Path to config file.')159parser.add_argument('input_path', type=str, help='Path to input image or directory.')160parser.add_argument('--output_path', type=str, default='outputs/', help='Output directory.')161parser.add_argument('--model_ckpt_path', type=str, default="", help='Output directory.')162parser.add_argument('--diffusion_steps', type=int, default=100, help='Denoising Sampling steps.')163parser.add_argument('--seed', type=int, default=42, help='Random seed for sampling.')164parser.add_argument('--scale', type=float, default=1.0, help='Scale of generated object.')165parser.add_argument('--materials', type=str_to_tuple, default=(1.0, 0.1), help=' metallic and roughness')166parser.add_argument('--distance', type=float, default=4.5, help='Render distance.')167parser.add_argument('--fov', type=float, default=30, help='Render distance.')168parser.add_argument('--env_path', type=str, default='data/env_mipmap/2', help='environment map')169parser.add_argument('--view', type=int, default=6, choices=[4, 6], help='Number of input views.')170parser.add_argument('--no_rembg', action='store_true', help='Do not remove input background.')171parser.add_argument('--export_texmap', action='store_true', help='Export a mesh with texture map.')172parser.add_argument('--save_video', action='store_true', help='Save a circular-view video.')173args = parser.parse_args()174seed_everything(args.seed)175 176###############################################################################177# Stage 0: Configuration.178###############################################################################179 180config = OmegaConf.load(args.config)181config_name = os.path.basename(args.config).replace('.yaml', '')182model_config = config.model_config183infer_config = config.infer_config184 185IS_FLEXICUBES = True186 187device = torch.device('cuda')188 189# load diffusion model190print('Loading diffusion model ...')191pipeline = DiffusionPipeline.from_pretrained(192 "sudo-ai/zero123plus-v1.2", 193 custom_pipeline="zero123plus",194 torch_dtype=torch.float16,195)196pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(197 pipeline.scheduler.config, timestep_spacing='trailing'198)199 200# load custom white-background UNet201print('Loading custom white-background unet ...')202if os.path.exists(infer_config.unet_path):203 unet_ckpt_path = infer_config.unet_path204else:205 unet_ckpt_path = hf_hub_download(repo_id="LTT/PRM", filename="diffusion_pytorch_model.bin", repo_type="model")206state_dict = torch.load(unet_ckpt_path, map_location='cpu')207pipeline.unet.load_state_dict(state_dict, strict=True)208 209pipeline = pipeline.to(device)210 211# load reconstruction model212print('Loading reconstruction model ...')213model = instantiate_from_config(model_config)214if os.path.exists(infer_config.model_path):215 model_ckpt_path = infer_config.model_path216else:217 model_ckpt_path = hf_hub_download(repo_id="LTT/PRM", filename="final_ckpt.ckpt", repo_type="model")218state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict']219state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')}220model.load_state_dict(state_dict, strict=True)221 222model = model.to(device)223if IS_FLEXICUBES:224 model.init_flexicubes_geometry(device, fovy=50.0)225model = model.eval()226 227# make output directories228image_path = os.path.join(args.output_path, config_name, 'images')229mesh_path = os.path.join(args.output_path, config_name, 'meshes')230video_path = os.path.join(args.output_path, config_name, 'videos')231os.makedirs(image_path, exist_ok=True)232os.makedirs(mesh_path, exist_ok=True)233os.makedirs(video_path, exist_ok=True)234 235# process input files236if os.path.isdir(args.input_path):237 input_files = [238 os.path.join(args.input_path, file) 239 for file in os.listdir(args.input_path) 240 if file.endswith('.png') or file.endswith('.jpg') or file.endswith('.webp')241 ]242else:243 input_files = [args.input_path]244print(f'Total number of input images: {len(input_files)}')245 246###############################################################################247# Stage 1: Multiview generation.248###############################################################################249 250rembg_session = None if args.no_rembg else rembg.new_session()251 252outputs = []253for idx, image_file in enumerate(input_files):254 name = os.path.basename(image_file).split('.')[0]255 print(f'[{idx+1}/{len(input_files)}] Imagining {name} ...')256 257 # remove background optionally258 input_image = Image.open(image_file)259 if not args.no_rembg:260 input_image = remove_background(input_image, rembg_session)261 input_image = resize_foreground(input_image, 0.85)262 # sampling263 output_image = pipeline(264 input_image, 265 num_inference_steps=args.diffusion_steps, 266 ).images[0]267 print(f"Image saved to {os.path.join(image_path, f'{name}.png')}")268 269 images = np.asarray(output_image, dtype=np.float32) / 255.0270 images = torch.from_numpy(images).permute(2, 0, 1).contiguous().float() # (3, 960, 640)271 images = rearrange(images, 'c (n h) (m w) -> (n m) c h w', n=3, m=2) # (6, 3, 320, 320)272 torchvision.utils.save_image(images, os.path.join(image_path, f'{name}.png'))273 sample = {'name': name, 'images': images}274 275# delete pipeline to save memory276# del pipeline277 278###############################################################################279# Stage 2: Reconstruction.280###############################################################################281 282 input_cameras = get_zero123plus_input_cameras(batch_size=1, radius=3.2*args.scale, fov=30).to(device)283 chunk_size = 20 if IS_FLEXICUBES else 1284 285# for idx, sample in enumerate(outputs):286 name = sample['name']287 print(f'[{idx+1}/{len(outputs)}] Creating {name} ...')288 289 images = sample['images'].unsqueeze(0).to(device)290 images = v2.functional.resize(images, 512, interpolation=3, antialias=True).clamp(0, 1)291 292 with torch.no_grad():293 # get triplane294 planes = model.forward_planes(images, input_cameras)295 296 mesh_path_idx = os.path.join(mesh_path, f'{name}.obj')297 298 mesh_out = model.extract_mesh(299 planes,300 use_texture_map=args.export_texmap,301 **infer_config,302 )303 if args.export_texmap:304 vertices, faces, uvs, mesh_tex_idx, tex_map = mesh_out305 save_obj_with_mtl(306 vertices.data.cpu().numpy(),307 uvs.data.cpu().numpy(),308 faces.data.cpu().numpy(),309 mesh_tex_idx.data.cpu().numpy(),310 tex_map.permute(1, 2, 0).data.cpu().numpy(),311 mesh_path_idx,312 )313 else:314 vertices, faces, vertex_colors = mesh_out315 save_obj(vertices, faces, vertex_colors, mesh_path_idx)316 print(f"Mesh saved to {mesh_path_idx}")317 318 render_size = 512319 if args.save_video:320 video_path_idx = os.path.join(video_path, f'{name}.mp4')321 render_size = infer_config.render_resolution322 ENV = load_mipmap(args.env_path)323 materials = args.materials324 325 all_mv, all_mvp, all_campos = get_render_cameras(326 batch_size=1, 327 M=240, 328 radius=args.distance, 329 elevation=(90, 60.0),330 is_flexicubes=IS_FLEXICUBES,331 fov=args.fov332 )333 334 frames, albedos, pbr_spec_lights, pbr_diffuse_lights, normals, alphas = render_frames(335 model, 336 planes, 337 render_cameras=all_mvp,338 camera_pos=all_campos,339 env=ENV,340 materials=materials,341 render_size=render_size, 342 chunk_size=chunk_size, 343 is_flexicubes=IS_FLEXICUBES,344 )345 normals = (torch.nn.functional.normalize(normals) + 1) / 2346 normals = normals * alphas + (1-alphas)347 all_frames = torch.cat([frames, albedos, pbr_spec_lights, pbr_diffuse_lights, normals], dim=3)348 349 # breakpoint()350 save_video(351 all_frames,352 video_path_idx,353 fps=30,354 )355 print(f"Video saved to {video_path_idx}")356 357