import pytorch_lightning as pl import os import sys import cv2 import time import json import torch import mcubes import trimesh import datetime import argparse import subprocess import numpy as np import gradio as gr from tqdm import tqdm import imageio.v2 as imageio from omegaconf import OmegaConf from safetensors.torch import load_file from huggingface_hub import hf_hub_download sys.path.append("3DTopia") from ldm.models.diffusion.ddim import DDIMSampler from ldm.models.diffusion.plms import PLMSSampler from ldm.models.diffusion.dpm_solver import DPMSolverSampler from utility.initialize import instantiate_from_config, get_obj_from_str from utility.triplane_renderer.eg3d_renderer import sample_from_planes, generate_planes from utility.triplane_renderer.renderer import get_rays, to8b import warnings warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=DeprecationWarning) def add_text(rgb, caption): font = cv2.FONT_HERSHEY_SIMPLEX # org gap = 10 org = (gap, gap) # fontScale fontScale = 0.3 # Blue color in BGR color = (255, 0, 0) # Line thickness of 2 px thickness = 1 break_caption = [] for i in range(len(caption) // 30 + 1): break_caption_i = caption[i*30:(i+1)*30] break_caption.append(break_caption_i) for i, bci in enumerate(break_caption): cv2.putText(rgb, bci, (gap, gap*(i+1)), font, fontScale, color, thickness, cv2.LINE_AA) return rgb config = "configs/default.yaml" local_ckpt = "checkpoints/3dtopia_diffusion_state_dict.ckpt" if os.path.exists(local_ckpt): ckpt = local_ckpt else: ckpt = hf_hub_download(repo_id="hongfz16/3DTopia", filename="model.safetensors") configs = OmegaConf.load(config) os.makedirs("tmp", exist_ok=True) if ckpt.endswith(".ckpt"): model = get_obj_from_str(configs.model["target"]).load_from_checkpoint(ckpt, map_location='cpu', strict=False, **configs.model.params) elif ckpt.endswith(".safetensors"): model = get_obj_from_str(configs.model["target"])(**configs.model.params) model_ckpt = load_file(ckpt) model.load_state_dict(model_ckpt) else: raise NotImplementedError device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") model = model.to(device) sampler = DDIMSampler(model) img_size = configs.model.params.unet_config.params.image_size channels = configs.model.params.unet_config.params.in_channels shape = [channels, img_size, img_size * 3] pose_folder = 'assets/sample_data/pose' poses_fname = sorted([os.path.join(pose_folder, f) for f in os.listdir(pose_folder)]) batch_rays_list = [] H = 128 ratio = 512 // H for p in poses_fname: c2w = np.loadtxt(p).reshape(4, 4) c2w[:3, 3] *= 2.2 c2w = np.array([ [1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1] ]) @ c2w k = np.array([ [560 / ratio, 0, H * 0.5], [0, 560 / ratio, H * 0.5], [0, 0, 1] ]) rays_o, rays_d = get_rays(H, H, torch.Tensor(k), torch.Tensor(c2w[:3, :4])) coords = torch.stack(torch.meshgrid(torch.linspace(0, H-1, H), torch.linspace(0, H-1, H), indexing='ij'), -1) coords = torch.reshape(coords, [-1,2]).long() rays_o = rays_o[coords[:, 0], coords[:, 1]] rays_d = rays_d[coords[:, 0], coords[:, 1]] batch_rays = torch.stack([rays_o, rays_d], 0) batch_rays_list.append(batch_rays) batch_rays_list = torch.stack(batch_rays_list, 0) def marching_cube(b, text, global_info): # prepare volumn for marching cube res = 128 assert 'decode_res' in global_info decode_res = global_info['decode_res'] c_list = torch.linspace(-1.2, 1.2, steps=res) grid_x, grid_y, grid_z = torch.meshgrid( c_list, c_list, c_list, indexing='ij' ) coords = torch.stack([grid_x, grid_y, grid_z], -1).to(device) plane_axes = generate_planes() feats = sample_from_planes( plane_axes, decode_res[b:b+1].reshape(1, 3, -1, 256, 256), coords.reshape(1, -1, 3), padding_mode='zeros', box_warp=2.4 ) fake_dirs = torch.zeros_like(coords) fake_dirs[..., 0] = 1 out = model.first_stage_model.triplane_decoder.decoder(feats, fake_dirs) u = out['sigma'].reshape(res, res, res).detach().cpu().numpy() del out # marching cube vertices, triangles = mcubes.marching_cubes(u, 10) min_bound = np.array([-1.2, -1.2, -1.2]) max_bound = np.array([1.2, 1.2, 1.2]) vertices = vertices / (res - 1) * (max_bound - min_bound)[None, :] + min_bound[None, :] pt_vertices = torch.from_numpy(vertices).to(device) # extract vertices color res_triplane = 256 render_kwargs = { 'depth_resolution': 128, 'disparity_space_sampling': False, 'box_warp': 2.4, 'depth_resolution_importance': 128, 'clamp_mode': 'softplus', 'white_back': True, 'det': True } rays_o_list = [ np.array([0, 0, 2]), np.array([0, 0, -2]), np.array([0, 2, 0]), np.array([0, -2, 0]), np.array([2, 0, 0]), np.array([-2, 0, 0]), ] rgb_final = None diff_final = None for rays_o in tqdm(rays_o_list): rays_o = torch.from_numpy(rays_o.reshape(1, 3)).repeat(vertices.shape[0], 1).float().to(device) rays_d = pt_vertices.reshape(-1, 3) - rays_o rays_d = rays_d / torch.norm(rays_d, dim=-1).reshape(-1, 1) dist = torch.norm(pt_vertices.reshape(-1, 3) - rays_o, dim=-1).cpu().numpy().reshape(-1) render_out = model.first_stage_model.triplane_decoder( decode_res[b:b+1].reshape(1, 3, -1, res_triplane, res_triplane), rays_o.unsqueeze(0), rays_d.unsqueeze(0), render_kwargs, whole_img=False, tvloss=False ) rgb = render_out['rgb_marched'].reshape(-1, 3).detach().cpu().numpy() depth = render_out['depth_final'].reshape(-1).detach().cpu().numpy() depth_diff = np.abs(dist - depth) if rgb_final is None: rgb_final = rgb.copy() diff_final = depth_diff.copy() else: ind = diff_final > depth_diff rgb_final[ind] = rgb[ind] diff_final[ind] = depth_diff[ind] # bgr to rgb rgb_final = np.stack([ rgb_final[:, 2], rgb_final[:, 1], rgb_final[:, 0] ], -1) # export to ply mesh = trimesh.Trimesh(vertices, triangles, vertex_colors=(rgb_final * 255).astype(np.uint8)) path = os.path.join('tmp', f"{text.replace(' ', '_')}_{str(datetime.datetime.now()).replace(' ', '_')}.ply") trimesh.exchange.export.export_mesh(mesh, path, file_type='ply') del vertices, triangles, rgb_final torch.cuda.empty_cache() return path def infer(prompt, samples, steps, scale, seed, global_info): prompt = prompt.replace('/', '') pl.seed_everything(seed) batch_size = samples with torch.no_grad(): noise = None c = model.get_learned_conditioning([prompt]) unconditional_c = torch.zeros_like(c) sample, _ = sampler.sample( S=steps, batch_size=batch_size, shape=shape, verbose=False, x_T = noise, conditioning = c.repeat(batch_size, 1, 1), unconditional_guidance_scale=scale, unconditional_conditioning=unconditional_c.repeat(batch_size, 1, 1) ) decode_res = model.decode_first_stage(sample) big_video_list = [] global_info['decode_res'] = decode_res for b in range(batch_size): def render_img(v): rgb_sample, _ = model.first_stage_model.render_triplane_eg3d_decoder( decode_res[b:b+1], batch_rays_list[v:v+1].to(device), torch.zeros(1, H, H, 3).to(device), ) rgb_sample = to8b(rgb_sample.detach().cpu().numpy())[0] rgb_sample = np.stack( [rgb_sample[..., 2], rgb_sample[..., 1], rgb_sample[..., 0]], -1 ) rgb_sample = add_text(rgb_sample, str(b)) return rgb_sample view_num = len(batch_rays_list) video_list = [] for v in tqdm(range(view_num//8*3, view_num//8*5, 2)): rgb_sample = render_img(v) video_list.append(rgb_sample) big_video_list.append(video_list) # if batch_size == 2: # cat_video_list = [ # np.concatenate([big_video_list[j][i] for j in range(len(big_video_list))], 1) \ # for i in range(len(big_video_list[0])) # ] # elif batch_size > 2: # if batch_size == 3: # big_video_list.append( # [np.zeros_like(f) for f in big_video_list[0]] # ) # cat_video_list = [ # np.concatenate([ # np.concatenate([big_video_list[0][i], big_video_list[1][i]], 1), # np.concatenate([big_video_list[2][i], big_video_list[3][i]], 1), # ], 0) \ # for i in range(len(big_video_list[0])) # ] # else: # cat_video_list = big_video_list[0] for _ in range(4 - batch_size): big_video_list.append( [np.zeros_like(f) + 255 for f in big_video_list[0]] ) cat_video_list = [ np.concatenate([ np.concatenate([big_video_list[0][i], big_video_list[1][i]], 1), np.concatenate([big_video_list[2][i], big_video_list[3][i]], 1), ], 0) \ for i in range(len(big_video_list[0])) ] path = f"tmp/{prompt.replace(' ', '_')}_{str(datetime.datetime.now()).replace(' ', '_')}.mp4" imageio.mimwrite(path, np.stack(cat_video_list, 0)) return global_info, path def infer_stage2(prompt, selection, seed, global_info): prompt = prompt.replace('/', '') mesh_path = marching_cube(int(selection), prompt, global_info) mesh_name = mesh_path.split('/')[-1][:-4] if2_cmd = f"threefiner if2 --mesh {mesh_path} --prompt \"{prompt}\" --outdir tmp --save {mesh_name}_if2.glb --text_dir --front_dir=-y" print(if2_cmd) # os.system(if2_cmd) subprocess.Popen(if2_cmd, shell=True).wait() torch.cuda.empty_cache() video_path = f"tmp/{prompt.replace(' ', '_')}_{str(datetime.datetime.now()).replace(' ', '_')}.mp4" render_cmd = f"kire {os.path.join('tmp', mesh_name + '_if2.glb')} --save_video {video_path} --wogui --force_cuda_rast --H 256 --W 256" print(render_cmd) # os.system(render_cmd) subprocess.Popen(render_cmd, shell=True).wait() torch.cuda.empty_cache() return video_path, os.path.join('tmp', mesh_name + '_if2.glb') markdown=f''' # 3DTopia A two-stage text-to-3D generation model. The first stage uses diffusion model to quickly generate candidates. The second stage refines the assets chosen from the first stage. ### Usage: First enter prompt for a 3D object, hit "Generate 3D". Then choose one candidate from the dropdown options for the second stage refinement and hit "Start Refinement". The final mesh can be downloaded from the bottom right box. ### Runtime: The first stage takes 30s if generating 4 samples. The second stage takes roughly 3 min. ### Useful links: [Github Repo](https://github.com/3DTopia/3DTopia) ''' block = gr.Blocks() with block: global_info = gr.State(dict()) gr.Markdown(markdown) with gr.Row(): with gr.Column(): with gr.Row(): text = gr.Textbox( label = "Enter your prompt", max_lines = 1, placeholder = "Enter your prompt", container = False, ) btn = gr.Button("Generate 3D") gallery = gr.Video(height=512) # advanced_button = gr.Button("Advanced Options", elem_id="advanced-btn") with gr.Row(elem_id="advanced-options"): with gr.Tab("Advanced options"): samples = gr.Slider(label="Number of Samples", minimum=1, maximum=4, value=4, step=1) steps = gr.Slider(label="Steps", minimum=1, maximum=500, value=50, step=1) scale = gr.Slider( label="Guidance Scale", minimum=0, maximum=50, value=7.5, step=0.1 ) seed = gr.Slider( label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True, ) gr.on([text.submit, btn.click], infer, inputs=[text, samples, steps, scale, seed, global_info], outputs=[global_info, gallery]) # advanced_button.click( # None, # [], # text, # ) with gr.Column(): with gr.Row(): dropdown = gr.Dropdown( ['0', '1', '2', '3'], label="Choose a candidate for stage2", value='0' ) btn_stage2 = gr.Button("Start Refinement") gallery = gr.Video(height=512) download = gr.File(label="Download mesh", file_count="single", height=100) gr.on([btn_stage2.click], infer_stage2, inputs=[text, dropdown, seed, global_info], outputs=[gallery, download]) block.launch(share=True)