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

def get_engine(use_fast):
    """Lazily build and cache the requested engine.

    EngineFast captures a CUDA graph for the doubled (cond+uncond) batch and so
    cannot run with guidance_scale == 0 (CFG disabled); Engine is used in that case.
    Both are cached separately so switching between them doesn't reload weights.
    """
    key = "engine_fast" if use_fast else "engine"
    if key 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_cls = EngineFast if use_fast else Engine
        GLOBAL_STATE[key] = engine_cls(
            config_path,
            gpt_ckpt_path,
            shape_ckpt_path,
            device=torch.device("cuda"),
        )
    return GLOBAL_STATE[key]

@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, guidance_scale=3.0):
    # 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,
        "guidance_scale": guidance_scale
    }
    
    # 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}, guidance_scale: {guidance_scale}")

    # EngineFast captures its CUDA graph for the doubled (cond+uncond) batch, so it
    # cannot run with guidance_scale == 0 (CFG disabled). Fall back to Engine there.
    use_fast = guidance_scale > 0.0
    engine = get_engine(use_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 = engine.t2s([input_prompt], use_kv_cache=True, resolution_base=resolution_base, bounding_box_xyz=normalized_bbox, guidance_scale=guidance_scale)
    # 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}
            **Disclaimer:** Content generated through the Hugging Face integration is not moderated by Roblox safety systems. Use of this integration is subject to the applicable license terms, and users are solely responsible for outputs generated.
            # 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 Bounding Box", 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)

                    guidance_scale = gr.Slider(
                        minimum=0.0,
                        maximum=10.0,
                        step=0.5,
                        value=3.0,
                        label="Guidance Scale",
                        info="Classifier-free guidance strength. Higher = stronger prompt adherence. 0 disables CFG (uses the slower non-fast engine).",
                    )
                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,
                guidance_scale
            ],
            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)