Upload code/cube3d/generate.py

code/cube3d/generate.py

@@ -0,0 +1,245 @@
+import argparse
+import os
+
+import torch
+import trimesh
+
+from cube3d.inference.engine import Engine, EngineFast
+from cube3d.inference.utils import normalize_bbox, select_device
+from cube3d.mesh_utils.postprocessing import (
+    PYMESHLAB_AVAILABLE,
+    create_pymeshset,
+    postprocess_mesh,
+    save_mesh,
+)
+from cube3d.renderer import renderer
+from cube3d.training.dataset import LegosTestDataset, LegosDataset
+from torch.utils.data.dataloader import DataLoader
+from cube3d.training.utils import normalize_bboxs
+
+def generate_mesh(
+    engine,
+    prompt,
+    output_dir,
+    output_name,
+    resolution_base=8.0,
+    disable_postprocess=False,
+    top_p=None,
+    bounding_box_xyz=None,
+):
+    #import ipdb; ipdb.set_trace()
+    mesh_v_f = engine.t2s(
+        [prompt],
+        use_kv_cache=True,
+        resolution_base=resolution_base,
+        top_p=top_p,
+        bounding_box_xyz=bounding_box_xyz,
+    )
+    vertices, faces = mesh_v_f[0][0], mesh_v_f[0][1]
+    obj_path = os.path.join(output_dir, f"{output_name}.obj")
+    if PYMESHLAB_AVAILABLE:
+        ms = create_pymeshset(vertices, faces)
+        if not disable_postprocess:
+            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, obj_path)
+
+        save_mesh(ms, obj_path)
+    else:
+        print(
+            "WARNING: pymeshlab is not available, using trimesh to export obj and skipping optional post processing."
+        )
+        mesh = trimesh.Trimesh(vertices, faces)
+        mesh.export(obj_path)
+
+    return obj_path
+
+def generate_ldr(
+    engine,
+    prompt,
+    inputs_ids,
+    output_dir,
+    output_name,
+    resolution_base=8.0,
+    disable_postprocess=False,
+    top_p=None,
+    bounding_box_xyz=None,
+    idx=None
+):
+    #import ipdb; ipdb.set_trace()
+    ldr = engine.t2l(
+        #[prompt],
+        prompt,
+        inputs_ids=inputs_ids,
+        use_kv_cache=True,
+        resolution_base=resolution_base,
+        top_p=top_p,
+        bounding_box_xyz=bounding_box_xyz,
+        idx=idx
+    )
+    # vertices, faces = mesh_v_f[0][0], mesh_v_f[0][1]
+    # obj_path = os.path.join(output_dir, f"{output_name}.obj")
+    # if PYMESHLAB_AVAILABLE:
+    #     ms = create_pymeshset(vertices, faces)
+    #     if not disable_postprocess:
+    #         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, obj_path)
+
+    #     save_mesh(ms, obj_path)
+    # else:
+    #     print(
+    #         "WARNING: pymeshlab is not available, using trimesh to export obj and skipping optional post processing."
+    #     )
+    #     mesh = trimesh.Trimesh(vertices, faces)
+    #     mesh.export(obj_path)
+
+    return ldr
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="cube shape generation script")
+    parser.add_argument(
+        "--config-path",
+        type=str,
+        default="cube3d/configs/open_model_v0.5.yaml",
+        help="Path to the configuration YAML file.",
+    )
+    parser.add_argument(
+        "--data-dir",
+        type=str,
+        required=True,
+        help="Path to the input dataset file.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="outputs/",
+        help="Path to the output directory to store .obj and .gif files",
+    )
+    parser.add_argument(
+        "--gpt-ckpt-path",
+        type=str,
+        required=True,
+        help="Path to the main GPT checkpoint file.",
+    )
+    parser.add_argument(
+        "--shape-ckpt-path",
+        type=str,
+        required=True,
+        help="Path to the shape encoder/decoder checkpoint file.",
+    )
+    parser.add_argument(
+        "--save-gpt-ckpt-path",
+        type=str,
+        required=True,
+        help="Path to the save adaption GPT checkpoint file.",
+    ) 
+    parser.add_argument(
+        "--fast-inference",
+        help="Use optimized inference",
+        default=False,
+        action="store_true",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Text prompt for generating a 3D mesh",
+    )
+    parser.add_argument(
+        "--top-p",
+        type=float,
+        default=None,
+        help="Float < 1: Keep smallest set of tokens with cumulative probability ≥ top_p. Default None: deterministic generation.",
+    )
+    parser.add_argument(
+        "--bounding-box-xyz",
+        nargs=3,
+        type=float,
+        help="Three float values for x, y, z bounding box",
+        default=None,
+        required=False,
+    )
+    parser.add_argument(
+        "--render-gif",
+        help="Render a turntable gif of the mesh",
+        default=False,
+        action="store_true",
+    )
+    parser.add_argument(
+        "--disable-postprocessing",
+        help="Disable postprocessing on the mesh. This will result in a mesh with more faces.",
+        default=False,
+        action="store_true",
+    )
+    parser.add_argument(
+        "--resolution-base",
+        type=float,
+        default=8.0,
+        help="Resolution base for the shape decoder.",
+    )
+    args = parser.parse_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+    device = select_device()
+    print(f"Using device: {device}")
+    # Initialize engine based on fast_inference flag
+    if args.fast_inference:
+        print(
+            "Using cuda graphs, this will take some time to warmup and capture the graph."
+        )
+        engine = EngineFast(
+            args.config_path, args.gpt_ckpt_path, args.shape_ckpt_path, device=device
+        )
+        print("Compiled the graph.")
+    else:
+        engine = Engine(
+            args.config_path, args.gpt_ckpt_path, args.shape_ckpt_path, args.save_gpt_ckpt_path, device=device
+        )
+
+    if args.bounding_box_xyz is not None:
+        args.bounding_box_xyz = normalize_bbox(tuple(args.bounding_box_xyz))
+
+    # Generate meshes based on input source
+    # obj_path = generate_mesh(
+    #     engine,
+    #     args.prompt,
+    #     args.output_dir,
+    #     "output",
+    #     args.resolution_base,
+    #     args.disable_postprocessing,
+    #     args.top_p,
+    #     args.bounding_box_xyz,
+    # )
+
+    test_dataset = LegosDataset(args)
+    batch_size = 1       
+    x_num = 213
+    y_num = 217
+    z_num = 529
+
+    # setup the dataloader
+    data_loader = DataLoader(
+        test_dataset,
+        shuffle=False,
+        batch_size=batch_size,
+    )
+    data_iter = iter(data_loader)
+    for idx in range(len(test_dataset)):
+        batch = next(data_iter)
+        prompt, targets, box = batch['prompt'], batch['target'].to(device), batch['bbox']
+        ldr = generate_ldr(
+            engine, 
+            prompt,
+            targets,
+            args.output_dir,
+            "output",
+            args.resolution_base,
+            args.disable_postprocessing,
+            args.top_p,
+            #args.bounding_box_xyz,
+            normalize_bboxs(box.float(), [x_num-1, y_num-1, z_num-1]), 
+            idx
+        )
+        # if idx>4:
+        #     break