add text2shape

app.py

@@ -1,7 +1,124 @@
+import argparse
 import gradio as gr
+import os
+import torch
+import trimesh
+import sys
+sys.path.append("cube")
+from cube3d.inference.engine import EngineFast
+from pathlib import Path
+import uuid
+import shutil
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+GLOBAL_STATE = {}
+
+def gen_save_folder(max_size=200):
+    os.makedirs(GLOBAL_STATE["SAVE_DIR"], exist_ok=True)
+
+    dirs = [f for f in Path(GLOBAL_STATE["SAVE_DIR"]).iterdir() if f.is_dir()]
+
+    if len(dirs) >= max_size:
+        oldest_dir = min(dirs, key=lambda x: x.stat().st_ctime)
+        shutil.rmtree(oldest_dir)
+        print(f"Removed the oldest folder: {oldest_dir}")
+
+    new_folder = os.path.join(GLOBAL_STATE["SAVE_DIR"], str(uuid.uuid4()))
+    os.makedirs(new_folder, exist_ok=True)
+    print(f"Created new folder: {new_folder}")
+
+    return new_folder
+
+def handle_text_prompt(input_prompt, ):
+    mesh_v_f = GLOBAL_STATE["engine_fast"].t2s([input_prompt], use_kv_cache=True, resolution_base=8.0)
+    # save output
+    vertices, faces = mesh_v_f[0][0], mesh_v_f[0][1]
+    save_folder = gen_save_folder()
+    output_path = os.path.join(save_folder, "output.obj")
+    trimesh.Trimesh(vertices=vertices, faces=faces).export(output_path)
+    return output_path
+
+def build_interface():
+    """Build UI for gradio app
+
+    Three tabs:
+    - Scene Generation
+    - Scene Completion
+    - Scene Understanding
+    """
+    title = "Cube 3D"
+    with gr.Blocks(theme=gr.themes.Soft(), title=title, fill_width=True) as interface:
+        gr.Markdown(
+            f"""
+            # {title}
+            """
+        )
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                with gr.Group():
+                    input_text_box = gr.Textbox(
+                        value=None,
+                        label="Prompt",
+                        lines=2,
+                    )
+                with gr.Row():
+                    submit_button = gr.Button("Submit", variant="primary")
+            with gr.Column(scale=3):
+                model3d = gr.Model3D(
+                    label="Output", height="45em", interactive=False
+                )
+    
+        submit_button.click(
+            handle_text_prompt,
+            inputs=[
+                input_text_box
+            ],
+            outputs=[
+                model3d
+            ]
+        )
+                
+    return interface
+
+if __name__=="__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config_path",
+        type=str,
+        help="Path to the config file",
+        default="cube/cube3d/configs/open_model.yaml",
+    )
+    parser.add_argument(
+        "--gpt_ckpt_path",
+        type=str,
+        help="Path to the gpt ckpt path",
+        default="model_weights/shape_gpt.safetensors",
+    )
+    parser.add_argument(
+        "--shape_ckpt_path",
+        type=str,
+        help="Path to the shape ckpt path",
+        default="model_weights/shape_tokenizer.safetensors",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="gradio_save_dir",
+    )
+
+    args = parser.parse_args()
+    engine_fast = EngineFast(
+        args.config_path, 
+        args.gpt_ckpt_path, 
+        args.shape_ckpt_path,
+        device=torch.device("cuda"),
+    )
+    GLOBAL_STATE["engine_fast"] = engine_fast
+    GLOBAL_STATE["SAVE_DIR"] = args.save_dir
+    os.makedirs(GLOBAL_STATE["SAVE_DIR"], exist_ok=True)
+
+    demo = build_interface()
+    demo.queue(default_concurrency_limit=None)
+    demo.launch()