app.py

import os
import sys
import cv2
import time
import tyro
import json
import kiui
import tqdm
import torch
import mcubes
import trimesh
import datetime
import argparse
import subprocess
import numpy as np
import gradio as gr
import imageio.v2 as imageio
import pytorch_lightning as pl
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

from threefiner.gui import GUI
from threefiner.opt import config_defaults, config_doc, check_options, Options

import warnings
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)

###################################### INIT STAGE 1 #########################################
config = "3DTopia/configs/default.yaml"
download_ckpt = "3DTopia/checkpoints/3dtopia_diffusion_state_dict.ckpt"
if not os.path.exists(download_ckpt):
    ckpt = hf_hub_download(repo_id="hongfz16/3DTopia", filename="model.safetensors")
else:
    ckpt = download_ckpt
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 = '3DTopia/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)
###################################### INIT STAGE 1 #########################################

###################################### INIT STAGE 2 #########################################
GRADIO_SAVE_PATH_MESH = 'gradio_output.glb'
GRADIO_SAVE_PATH_VIDEO = 'gradio_output.mp4'

# opt = tyro.cli(tyro.extras.subcommand_type_from_defaults(config_defaults, config_doc))
opt = Options(
    mode='IF2',
    iters=400,
)

# hacks for not loading mesh at initialization
# opt.mesh = 'tmp/_2024-01-25_19:33:03.110191_if2.glb'
opt.save = GRADIO_SAVE_PATH_MESH
opt.prompt = ''
opt.text_dir = True
opt.front_dir = '+z'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
gui = GUI(opt)
###################################### INIT STAGE 2 #########################################

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

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.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.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, iters):
    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)
    # 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)
    # subprocess.Popen(render_cmd, shell=True).wait()
    # torch.cuda.empty_cache()

    process_stage2(mesh_path, prompt, "down", iters, f'tmp/{mesh_name}_if2.glb', video_path)
    torch.cuda.empty_cache()

    return video_path, f'tmp/{mesh_name}_if2.glb'

def process_stage2(input_model, input_text, input_dir, iters, output_model, output_video):
    # set front facing direction (map from gradio model3D's mysterious coordinate system to OpenGL...)
    opt.text_dir = True
    if input_dir == 'front':
        opt.front_dir = '-z'
    elif input_dir == 'back':
        opt.front_dir = '+z'
    elif input_dir == 'left':
        opt.front_dir = '+x'
    elif input_dir == 'right':
        opt.front_dir = '-x'
    elif input_dir == 'up':
        opt.front_dir = '+y'
    elif input_dir == 'down':
        opt.front_dir = '-y'
    else:
        # turn off text_dir
        opt.text_dir = False
        opt.front_dir = '+z'
    
    # set mesh path
    opt.mesh = input_model

    # load mesh!
    gui.renderer = gui.renderer_class(opt, device).to(device)

    # set prompt
    gui.prompt = opt.positive_prompt + ', ' + input_text

    # train
    gui.prepare_train() # update optimizer and prompt embeddings
    for i in tqdm.trange(iters):
        gui.train_step()

    # save mesh & video
    gui.save_model(output_model)
    gui.save_model(output_video)

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 1m30s.

  ### 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)
            with gr.Row(elem_id="advanced-options"):
                with gr.Tab("Advanced options"):
                    # input_dir = gr.Radio(['front', 'back', 'left', 'right', 'up', 'down'], value='down', label="front-facing direction")
                    iters = gr.Slider(minimum=100, maximum=1000, step=100, value=400, label="Refine iterations")
            download = gr.File(label="Download Mesh", file_count="single", height=100)
            gr.on([btn_stage2.click], infer_stage2, inputs=[text, dropdown, seed, global_info, iters], outputs=[gallery, download])

block.launch(share=True)