support .glb export and update README

README.md

@@ -1,13 +1,29 @@
 ---
-title: One 2 3 45
-emoji: 📈
+title: One-2-3-45
+emoji: 📸🚀🌟
 colorFrom: red
 colorTo: yellow
 sdk: gradio
 sdk_version: 3.36.1
 app_file: app.py
-pinned: false
+pinned: true
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# One-2-3-45: Any Single Image to 3D Mesh in 45 Seconds without Per-Shape Optimization
+
+Paper: https://arxiv.org/abs/2306.16928
+Code: https://github.com/One-2-3-45/One-2-3-45
+
+## BibTeX
+
+```bibtex
+@misc{liu2023one2345,
+      title={One-2-3-45: Any Single Image to 3D Mesh in 45 Seconds without Per-Shape Optimization}, 
+      author={Minghua Liu and Chao Xu and Haian Jin and Linghao Chen and Mukund Varma T and Zexiang Xu and Hao Su},
+      year={2023},
+      eprint={2306.16928},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

app.py

@@ -348,7 +348,7 @@ def stage1_run(models, device, cam_vis, tmp_dir,
         return (rerun_all, *reset, *outputs)
     
 def stage2_run(models, device, tmp_dir,
-               elev, scale, rerun_all=[], stage2_steps=50):
+               elev, scale, is_glb=False, rerun_all=[], stage2_steps=50):
     flag_lower_cam = 90-int(elev["label"]) <= 75
     is_rerun = True if rerun_all else False
     model = models['turncam'].half()
@@ -373,7 +373,8 @@ def stage2_run(models, device, tmp_dir,
     os.chdir(main_dir_path)
 
     ply_path = os.path.join(tmp_dir, f"meshes_val_bg/lod0/mesh_00215000_gradio_lod0.ply")
-    mesh_path = os.path.join(tmp_dir, "mesh.obj")
+    mesh_ext = ".glb" if is_glb else ".obj"
+    mesh_path = os.path.join(tmp_dir, f"mesh{mesh_ext}")
     # Read the textured mesh from .ply file
     mesh = trimesh.load_mesh(ply_path)
     axis = [1, 0, 0]
@@ -388,7 +389,10 @@ def stage2_run(models, device, tmp_dir,
     mesh.vertices[:, 0] = -mesh.vertices[:, 0]
     mesh.faces = np.fliplr(mesh.faces)
     # Export the mesh as .obj file with colors
-    mesh.export(mesh_path, file_type='obj', include_color=True)
+    if not is_glb:
+        mesh.export(mesh_path, file_type='obj', include_color=True)
+    else:
+        mesh.export(mesh_path, file_type='glb')
     torch.cuda.empty_cache()
 
     if not is_rerun:
@@ -523,6 +527,8 @@ def run_demo(
                                              label='Diffusion guidance scale')
                     steps_slider = gr.Slider(5, 200, value=75, step=5,
                                              label='Number of diffusion inference steps')
+                    glb_chk = gr.Checkbox(
+                        False, label='Export the mesh in .glb format')
 
                 run_btn = gr.Button('Run Generation', variant='primary', interactive=False)
                 guide_text = gr.Markdown(_USER_GUIDE, visible=True)
@@ -642,7 +648,7 @@ def run_demo(
                                 outputs=[elev_output, vis_output, *views]
                       ).success(fn=partial(update_guide, _GEN_2), outputs=[guide_text], queue=False
                       ).success(fn=partial(stage2_run, models, device),
-                                inputs=[tmp_dir, elev_output, scale_slider],
+                                inputs=[tmp_dir, elev_output, scale_slider, glb_chk],
                                 outputs=[mesh_output]
                       ).success(fn=partial(update_guide, _DONE), outputs=[guide_text], queue=False)
     
@@ -652,7 +658,7 @@ def run_demo(
                              outputs=[rerun_idx, *btn_retrys, *views]
                             ).success(fn=partial(update_guide, _REGEN_1), outputs=[guide_text], queue=False)
         regen_mesh_btn.click(fn=partial(stage2_run, models, device),
-                             inputs=[tmp_dir, elev_output, scale_slider, rerun_idx],
+                             inputs=[tmp_dir, elev_output, scale_slider, glb_chk, rerun_idx],
                              outputs=[mesh_output, rerun_idx, regen_view_btn, regen_mesh_btn]
                             ).success(fn=partial(update_guide, _REGEN_2), outputs=[guide_text], queue=False)