bbox sliders always present

Dockerfile

@@ -30,6 +30,6 @@ RUN git clone https://github.com/Roblox/cube.git
 
 WORKDIR /home/user/app/cube
 RUN pip install .[meshlab]
-RUN huggingface-cli download Roblox/cube3d-v0.1 --local-dir ./model_weights
+RUN huggingface-cli download Roblox/cube3d-v0.5 --local-dir ./model_weights
 
 WORKDIR /home/user/app

app.py

@@ -6,6 +6,7 @@ import torch
 import trimesh
 import sys
 from pathlib import Path
+import numpy as np
 
 pathdir = Path(__file__).parent / 'cube'
 sys.path.append(pathdir.as_posix())
@@ -49,8 +50,8 @@ def gen_save_folder(max_size=200):
     return new_folder
 
 @spaces.GPU
-def handle_text_prompt(input_prompt, bbox_option="No Bounding Box", bbox_x=1.0, bbox_y=1.0, bbox_z=1.0, hi_res=False):
-    print(f"prompt: {input_prompt}, bbox_option: {bbox_option}, bbox_x: {bbox_x}, bbox_y: {bbox_y}, bbox_z: {bbox_z}, hi_res: {hi_res}")
+def handle_text_prompt(input_prompt, use_bbox = True, bbox_x=1.0, bbox_y=1.0, bbox_z=1.0, hi_res=False):
+    print(f"prompt: {input_prompt}, use_bbox: {use_bbox}, bbox_x: {bbox_x}, bbox_y: {bbox_y}, bbox_z: {bbox_z}, hi_res: {hi_res}")
 
     if "engine_fast" not in GLOBAL_STATE: 
         config_path = GLOBAL_STATE["config_path"]
@@ -66,19 +67,13 @@ def handle_text_prompt(input_prompt, bbox_option="No Bounding Box", bbox_x=1.0,
 
     # Determine bounding box size based on option
     bbox_size = None
-    if bbox_option == "1x1x1":
-        bbox_size = [1.0, 1.0, 1.0]
-    elif bbox_option == "1x2x1":
-        bbox_size = [1.0, 2.0, 1.0]
-    elif bbox_option == "2x3x1":
-        bbox_size = [2.0, 3.0, 1.0]
-    elif bbox_option == "Custom":
+    if use_bbox:
         bbox_size = [bbox_x, bbox_y, bbox_z]
     # For "No Bounding Box", bbox_size remains None
     
     normalized_bbox = normalize_bbox(bbox_size) if bbox_size is not None else None
 
-    resolution_base = 8.5 if hi_res else 8.0
+    resolution_base = 9.0 if hi_res else 8.0
     mesh_v_f = GLOBAL_STATE["engine_fast"].t2s([input_prompt], use_kv_cache=True, resolution_base=resolution_base, bounding_box_xyz=normalized_bbox)
     # save output
     vertices, faces = mesh_v_f[0][0], mesh_v_f[0][1]
@@ -91,9 +86,17 @@ def handle_text_prompt(input_prompt, bbox_option="No Bounding Box", bbox_x=1.0,
     vertices = mesh.vertex_matrix()
     faces = mesh.face_matrix()
 
+    min_extents = np.min(mesh.vertex_matrix(), axis = 0)
+    max_extents = np.max(mesh.vertex_matrix(), axis = 0)
+
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+    scene = trimesh.scene.Scene()
+    scene.add_geometry(mesh)
+
     save_folder = gen_save_folder()
     output_path = os.path.join(save_folder, "output.glb")
-    trimesh.Trimesh(vertices=vertices, faces=faces).export(output_path)
+    # trimesh.Trimesh(vertices=vertices, faces=faces).export(output_path)
+    scene.export(output_path)
     return output_path
 
 def build_interface():
@@ -116,25 +119,27 @@ def build_interface():
                         label="Prompt",
                         lines=2,
                     )
-                    bbox_option = gr.Dropdown(
-                        label="Bounding Box",
-                        choices=["No Bounding Box", "1x1x1", "1x2x1", "2x3x1", "Custom"],
-                        value="No Bounding Box"
+
+                    use_bbox = gr.Checkbox(label="Use Bbox", value=False)
+
+                    with gr.Group() as bbox_group:
+                        bbox_x = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Length", interactive=False)
+                        bbox_y = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Height", interactive=False)
+                        bbox_z = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Depth", interactive=False)
+
+                    # Enable/disable bbox sliders based on use_bbox checkbox
+                    def toggle_bbox_interactivity(use_bbox):
+                        return (
+                            gr.Slider(interactive=use_bbox),
+                            gr.Slider(interactive=use_bbox),
+                            gr.Slider(interactive=use_bbox)
+                        )
+                    use_bbox.change(
+                        toggle_bbox_interactivity,
+                        inputs=[use_bbox],
+                        outputs=[bbox_x, bbox_y, bbox_z]
                     )
-                    with gr.Group(visible=False) as bbox_group:
-                        bbox_x = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Bbox X")
-                        bbox_y = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Bbox Y")
-                        bbox_z = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Bbox Z")
                     
-                    # Show/hide bbox sliders based on dropdown
-                    def toggle_bbox_visibility(bbox_option):
-                        return gr.Group(visible=bbox_option == "Custom")
-                    
-                    bbox_option.change(
-                        toggle_bbox_visibility,
-                        inputs=[bbox_option],
-                        outputs=[bbox_group]
-                    )
                     hi_res = gr.Checkbox(label="Hi-Res", value=False)
                 with gr.Row():
                     submit_button = gr.Button("Submit", variant="primary")
@@ -147,7 +152,7 @@ def build_interface():
             handle_text_prompt,
             inputs=[
                 input_text_box,
-                bbox_option,
+                use_bbox,
                 bbox_x,
                 bbox_y,
                 bbox_z,
@@ -159,7 +164,6 @@ def build_interface():
         )
                 
     return interface
-
 def generate(args):
     GLOBAL_STATE["config_path"] = args.config_path
     GLOBAL_STATE["SAVE_DIR"] = args.save_dir