app.py

import spaces
import argparse
import gradio as gr
import os
import torch
import trimesh
import sys
from pathlib import Path
import numpy as np
import json
from datetime import datetime

pathdir = Path(__file__).parent / 'cube'
sys.path.append(pathdir.as_posix())

# print(__file__)
# print(os.listdir())
# print(os.listdir('cube'))
# print(pathdir.as_posix())

from cube3d.inference.engine import EngineFast, Engine
from cube3d.inference.utils import normalize_bbox
from pathlib import Path
import uuid
import shutil
from huggingface_hub import snapshot_download

from cube3d.mesh_utils.postprocessing import (
    PYMESHLAB_AVAILABLE,
    create_pymeshset,
    postprocess_mesh,
    save_mesh,
)


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

@spaces.GPU
def handle_text_prompt(input_prompt, use_bbox = True, bbox_x=1.0, bbox_y=1.0, bbox_z=1.0, hi_res=False):
    # Create debug info
    debug_info = {
        "timestamp": datetime.now().isoformat(),
        "prompt": input_prompt,
        "use_bbox": use_bbox,
        "bbox_x": bbox_x,
        "bbox_y": bbox_y,
        "bbox_z": bbox_z,
        "hi_res": hi_res
    }
    
    # Save to persistent storage
    data_dir = "/data"
    os.makedirs(data_dir, exist_ok=True)
    
    prompt_file = os.path.join(data_dir, "prompt_log.jsonl")
    with open(prompt_file, "a") as f:
        f.write(json.dumps(debug_info) + "\n")
    
    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"]
        gpt_ckpt_path = "./model_weights/shape_gpt.safetensors"
        shape_ckpt_path = "./model_weights/shape_tokenizer.safetensors"
        engine_fast = EngineFast(
            config_path,
            gpt_ckpt_path, 
            shape_ckpt_path,
            device=torch.device("cuda"),
        )
        GLOBAL_STATE["engine_fast"] = engine_fast

    # Determine bounding box size based on option
    bbox_size = None
    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 = 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]

    ms = create_pymeshset(vertices, faces)
    target_face_num = max(10000, int(faces.shape[0] * 0.1))
    print(f"Postprocessing mesh to {target_face_num} faces")
    postprocess_mesh(ms, target_face_num)
    mesh = ms.current_mesh()
    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)
    scene.export(output_path)
    return output_path

def build_interface():
    """Build UI for gradio app
    """
    title = "Cube 3D"
    with gr.Blocks(theme=gr.themes.Soft(), title=title, fill_width=True) as interface:
        gr.Markdown(
            f"""
            # {title}
            # Check out our [Github](https://github.com/Roblox/cube) to try it on your own machine!
            """
        )

        with gr.Row():
            with gr.Column(scale=2):
                with gr.Group():
                    input_text_box = gr.Textbox(
                        value=None,
                        label="Prompt",
                        lines=2,
                    )

                    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]
                    )
                    
                    hi_res = gr.Checkbox(label="Hi-Res", value=False)
                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,
                use_bbox,
                bbox_x,
                bbox_y,
                bbox_z,
                hi_res
            ],
            outputs=[
                model3d
            ]
        )
                
    return interface
def generate(args):
    GLOBAL_STATE["config_path"] = args.config_path
    GLOBAL_STATE["SAVE_DIR"] = args.save_dir
    os.makedirs(GLOBAL_STATE["SAVE_DIR"], exist_ok=True)

    demo = build_interface()
    demo.queue(default_concurrency_limit=1)
    demo.launch()

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_v0.5.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()
    snapshot_download(
        repo_id="Roblox/cube3d-v0.5",
        local_dir="./model_weights"
    )
    generate(args)