import spaces import os import imageio import numpy as np import torch import rembg from PIL import Image from torchvision.transforms import v2 from pytorch_lightning import seed_everything from omegaconf import OmegaConf from einops import rearrange, repeat from tqdm import tqdm import threading from queue import SimpleQueue from typing import Any from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler import rerun as rr from gradio_rerun import Rerun import src from src.utils.train_util import instantiate_from_config from src.utils.camera_util import ( FOV_to_intrinsics, get_zero123plus_input_cameras, get_circular_camera_poses, ) from src.utils.mesh_util import save_obj, save_glb from src.utils.infer_util import remove_background, resize_foreground, images_to_video from src.models.lrm_mesh import InstantMesh import tempfile from functools import partial from huggingface_hub import hf_hub_download import gradio as gr def get_render_cameras(batch_size=1, M=120, radius=2.5, elevation=10.0, is_flexicubes=False): """ Get the rendering camera parameters. """ c2ws = get_circular_camera_poses(M=M, radius=radius, elevation=elevation) if is_flexicubes: cameras = torch.linalg.inv(c2ws) cameras = cameras.unsqueeze(0).repeat(batch_size, 1, 1, 1) else: extrinsics = c2ws.flatten(-2) intrinsics = FOV_to_intrinsics(50.0).unsqueeze(0).repeat(M, 1, 1).float().flatten(-2) cameras = torch.cat([extrinsics, intrinsics], dim=-1) cameras = cameras.unsqueeze(0).repeat(batch_size, 1, 1) return cameras def images_to_video(images, output_path, fps=30): # images: (N, 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).clip(0, 255) 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) imageio.mimwrite(output_path, np.stack(frames), fps=fps, codec='h264') ############################################################################### # Configuration. ############################################################################### import shutil def find_cuda(): # Check if CUDA_HOME or CUDA_PATH environment variables are set cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') if cuda_home and os.path.exists(cuda_home): return cuda_home # Search for the nvcc executable in the system's PATH nvcc_path = shutil.which('nvcc') if nvcc_path: # Remove the 'bin/nvcc' part to get the CUDA installation path cuda_path = os.path.dirname(os.path.dirname(nvcc_path)) return cuda_path return None cuda_path = find_cuda() if cuda_path: print(f"CUDA installation found at: {cuda_path}") else: print("CUDA installation not found") config_path = 'configs/instant-mesh-large.yaml' config = OmegaConf.load(config_path) config_name = os.path.basename(config_path).replace('.yaml', '') model_config = config.model_config infer_config = config.infer_config IS_FLEXICUBES = True if config_name.startswith('instant-mesh') else False device = torch.device('cuda') # load diffusion model print('Loading diffusion model ...') pipeline = DiffusionPipeline.from_pretrained( "sudo-ai/zero123plus-v1.2", custom_pipeline="zero123plus", torch_dtype=torch.float16, ) pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( pipeline.scheduler.config, timestep_spacing='trailing' ) # load custom white-background UNet unet_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="diffusion_pytorch_model.bin", repo_type="model") state_dict = torch.load(unet_ckpt_path, map_location='cpu') pipeline.unet.load_state_dict(state_dict, strict=True) pipeline = pipeline.to(device) print(f'type(pipeline)={type(pipeline)}') # load reconstruction model print('Loading reconstruction model ...') model_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="instant_mesh_large.ckpt", repo_type="model") model: InstantMesh = instantiate_from_config(model_config) state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict'] state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.') and 'source_camera' not in k} model.load_state_dict(state_dict, strict=True) model = model.to(device) print('Loading Finished!') def check_input_image(input_image): if input_image is None: raise gr.Error("No image uploaded!") def preprocess(input_image, do_remove_background): rembg_session = rembg.new_session() if do_remove_background else None if do_remove_background: input_image = remove_background(input_image, rembg_session) input_image = resize_foreground(input_image, 0.85) return input_image def pipeline_callback(output_queue: SimpleQueue, pipe: Any, step_index: int, timestep: float, callback_kwargs: dict[str, Any]) -> dict[str, Any]: rr.set_time_sequence("iteration", step_index) rr.set_time_seconds("timestep", timestep) latents = callback_kwargs["latents"] image = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0] # type: ignore[attr-defined] image = pipe.image_processor.postprocess(image, output_type="np").squeeze() # type: ignore[attr-defined] output_queue.put(("log", "mvs/image", rr.Image(image))) output_queue.put(("log", "mvs/latents", rr.Tensor(latents.squeeze()))) return callback_kwargs @spaces.GPU def generate_mvs(input_image, sample_steps, sample_seed): seed_everything(sample_seed) def thread_target(output_queue, input_image, sample_steps): z123_image = pipeline( input_image, num_inference_steps=sample_steps, callback_on_step_end=lambda *args, **kwargs: pipeline_callback(output_queue, *args, **kwargs), ).images[0] output_queue.put(("z123_image", z123_image)) output_queue = SimpleQueue() z123_thread = threading.Thread( target=thread_target, args= [ output_queue, input_image, sample_steps, ] ) z123_thread.start() while True: msg = output_queue.get() yield msg if msg[0] == "z123_image": break z123_thread.join() @spaces.GPU def make3d(images: Image.Image): output_queue = SimpleQueue() handle = threading.Thread(target=_make3d, args=[output_queue, images]) handle.start() while True: msg = output_queue.get() yield msg if msg[0] == "mesh": break handle.join() def _make3d(output_queue: SimpleQueue, images: Image.Image): print(f'type(images)={type(images)}') global model if IS_FLEXICUBES: model.init_flexicubes_geometry(device, use_renderer=False) model = model.eval() images = np.asarray(images, dtype=np.float32) / 255.0 images = torch.from_numpy(images).permute(2, 0, 1).contiguous().float() # (3, 960, 640) images = rearrange(images, 'c (n h) (m w) -> (n m) c h w', n=3, m=2) # (6, 3, 320, 320) input_cameras = get_zero123plus_input_cameras(batch_size=1, radius=4.0).to(device) render_cameras = get_render_cameras(batch_size=1, radius=2.5, is_flexicubes=IS_FLEXICUBES).to(device) print(f'type(input_cameras)={type(input_cameras)}') images = images.unsqueeze(0).to(device) images = v2.functional.resize(images, (320, 320), interpolation=3, antialias=True).clamp(0, 1) print(f'type(images)={type(images)}') mesh_fpath = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False).name print(mesh_fpath) mesh_basename = os.path.basename(mesh_fpath).split('.')[0] mesh_dirname = os.path.dirname(mesh_fpath) video_fpath = os.path.join(mesh_dirname, f"{mesh_basename}.mp4") mesh_glb_fpath = os.path.join(mesh_dirname, f"{mesh_basename}.glb") with torch.no_grad(): # get triplane planes = model.forward_planes(images, input_cameras) print(f'type(planes)={type(planes)}') # # get video chunk_size = 20 if IS_FLEXICUBES else 1 render_size = 384 # frames = [] for i in tqdm(range(0, render_cameras.shape[1], chunk_size)): if IS_FLEXICUBES: frame = model.forward_geometry( planes, render_cameras[:, i:i+chunk_size], render_size=render_size, )['img'] else: frame = model.synthesizer( planes, cameras=render_cameras[:, i:i+chunk_size], render_size=render_size, )['images_rgb'] print(f'type(frame)={type(frame)}') output_queue.put(("log", "3dvideo", rr.Image(frame))) # frames.append(frame) # frames = torch.cat(frames, dim=1) # images_to_video( # frames[0], # video_fpath, # fps=30, # ) # print(f"Video saved to {video_fpath}") # get mesh mesh_out = model.extract_mesh( planes, use_texture_map=False, **infer_config, ) print(f'type(mesh_out)={type(mesh_out)}') vertices, faces, vertex_colors = mesh_out vertices = vertices[:, [1, 2, 0]] print(f'type(vertices)={type(vertices)}') print(f'type(faces)={type(faces)}') print(f'type(vertex_colors)={type(vertex_colors)}') save_glb(vertices, faces, vertex_colors, mesh_glb_fpath) save_obj(vertices, faces, vertex_colors, mesh_fpath) print(f"Mesh saved to {mesh_fpath}") output_queue.put(("mesh", mesh_out)) @rr.thread_local_stream("InstantMesh") def log_to_rr(input_image, do_remove_background, sample_steps, sample_seed): preprocessed_image = preprocess(input_image, do_remove_background) stream = rr.binary_stream() rr.log("preprocessed_image", rr.Image(preprocessed_image)) yield stream.read() for msg in generate_mvs(input_image, sample_steps, sample_seed): if msg[0] == "z123_image": z123_image = msg[1] break elif msg[0] == "log": entity_path = msg[1] entity = msg[2] rr.log(entity_path, entity) yield stream.read() rr.log("z123image", rr.Image(z123_image)) yield stream.read() for msg in make3d(z123_image): if msg[0] == "log": rr.log(msg[1], msg[2]) _HEADER_ = '''