3DTopia test

app.py

@@ -2,7 +2,10 @@ import os
 import sys
 import cv2
 import time
+import tyro
 import json
+import kiui
+import tqdm
 import torch
 import mcubes
 import trimesh
@@ -11,7 +14,6 @@ import argparse
 import subprocess
 import numpy as np
 import gradio as gr
-from tqdm import tqdm
 import imageio.v2 as imageio
 import pytorch_lightning as pl
 from omegaconf import OmegaConf
@@ -28,10 +30,90 @@ 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
@@ -51,76 +133,6 @@ def add_text(rgb, caption):
         cv2.putText(rgb, bci, (gap, gap*(i+1)), font, fontScale, color, thickness, cv2.LINE_AA)
     return rgb
 
-config = "3DTopia/configs/default.yaml"
-# local_ckpt = "3DTopia/checkpoints/3dtopia_diffusion_state_dict.ckpt"
-local_ckpt = "/data/3DTopia_all/3DTopia_code/checkpoints/model.safetensors"
-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)
-
-import sys
-import traceback
-
-try:
-    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)
-        print("download finish")
-        model_ckpt = load_file(ckpt)
-        print("download finish")
-        model.load_state_dict(model_ckpt)
-        print("download finish")
-    else:
-        raise NotImplementedError
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    model = model.to(device)
-    print("download finish")
-    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)
-except Exception as e:
-    print(e)
-    print(traceback.format_exc())
-    print(sys.exc_info()[2])
-
-
-print("download finish")
 def marching_cube(b, text, global_info):
     # prepare volumn for marching cube
     res = 128
@@ -169,7 +181,7 @@ def marching_cube(b, text, global_info):
     ]
     rgb_final = None
     diff_final = None
-    for rays_o in tqdm(rays_o_list):
+    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)
@@ -246,7 +258,7 @@ def infer(prompt, samples, steps, scale, seed, global_info):
 
             view_num = len(batch_rays_list)
             video_list = []
-            for v in tqdm(range(view_num//8*3, view_num//8*5, 2)):
+            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)
@@ -287,25 +299,62 @@ def infer(prompt, samples, steps, scale, seed, global_info):
 
     return global_info, path
 
-def infer_stage2(prompt, selection, seed, global_info):
+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)
-    # os.system(if2_cmd)
-    subprocess.Popen(if2_cmd, shell=True).wait()
-    torch.cuda.empty_cache()
-
+    # 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)
-    # os.system(render_cmd)
-    subprocess.Popen(render_cmd, shell=True).wait()
+    # 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, os.path.join('tmp', mesh_name + '_if2.glb')
+    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
@@ -315,7 +364,7 @@ markdown=f'''
   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.
+  The first stage takes 30s if generating 4 samples. The second stage takes roughly 1m30s.
 
   ### Useful links:
   [Github Repo](https://github.com/3DTopia/3DTopia)
@@ -337,7 +386,7 @@ with block:
                 )
                 btn = gr.Button("Generate 3D")
             gallery = gr.Video(height=512)
-            # advanced_button = gr.Button("Advanced Options", elem_id="advanced-btn")
+            # 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)
@@ -361,11 +410,15 @@ with block:
         with gr.Column():
             with gr.Row():
                 dropdown = gr.Dropdown(
-                    ['0', '1', '2', '3'], label="Choose a candidate for stage2", value='0'
+                    ['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])
+            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, debug=True)
+block.launch(share=True)

requirements.txt

@@ -54,4 +54,7 @@ trimesh
 vit-pytorch
 wandb
 wcwidth
-zipp
+zipp
+git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
+git+https://github.com/NVlabs/nvdiffrast
+git+https://github.com/3DTopia/threefiner