adding cube sources

cube/.gitignore

@@ -0,0 +1,176 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+.vscode/*
+
+.DS_Store
+
+# Output folder
+outputs/
+model_weights/

cube/LICENSE

@@ -0,0 +1,80 @@
+CUBE3D RESEARCH-ONLY RAIL-MS LICENSE
+
+Licensed Artifacts:
+Cube3d-v0.1 and related inference code 
+
+I. SCOPE
+This Research-Only RAIL License is generally applicable to the Artifacts identified above.
+For valuable consideration, You and Licensor agree as follows:
+1. Definitions
+(a) “Artifact” means a software application (in either binary or source code format), Model, or Source Code, in accordance with what are specified above as the “Licensed Artifacts.”
+(b) “Contribution” means any work, including any modifications or additions to an Artifact, that is intentionally submitted to Licensor for inclusion or incorporation in the Artifact directly or indirectly by the rights owner. For the purposes of this definition, “submitted” means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing, sharing and improving the Artifact, but excluding communication that is conspicuously marked or otherwise designated in writing by the contributor as “Not a Contribution.”
+(c) “Contributor” means Licensor or any other individual or legal entity that creates or owns a Contribution that is added to or incorporated into an Artifact or Derivative.
+(d) “Data” means a collection of information or content extracted from the dataset used with a given Model, including to train, pretrain, or otherwise evaluate the Model. 
+(e) “Derivative” means a work derived from or based upon an Artifact, and includes all modified versions of such Artifact.
+(f) “Distribution” means any transmission, reproduction, publication or other sharing of an Artifact or Derivative to a third party, including providing a hosted service incorporating the Artifact, which is made available by electronic or other remote means—e.g., API-based or web access.
+(g) “License” means the terms and conditions for use, reproduction, and Distribution as defined in this document.
+(h) “Licensor” means the rights owner (by virtue of creation or documented transfer of ownership) or entity authorized by the rights owner (e.g., exclusive licensee) that is granting the rights in this License.
+(i) “Model” means any machine-learning based assembly or assemblies (including checkpoints), consisting of learnt weights, parameters (including optimizer states), corresponding to the model architecture as embodied in the Source Code.
+(j) “Output” means the results of operating a Model as embodied in informational content resulting therefrom.
+(k) “Permitted Purpose” means for academic or research purposes only.
+(l) “Source Code” means any collection of text written using human-readable programming language, including the code and scripts used to define, run, load, benchmark or evaluate a Model or any component thereof, or used to prepare data for training or evaluation. Source Code includes any accompanying documentation, tutorials and examples. For clarity, the term “Source Code” as used in this License includes any and all Derivatives of such Source Code.
+(m) “Third Party” means any individual or legal entity that is not under common control with Licensor or You.
+(n) “Use,” with respect to an Artifact, means accessing, using, copying, modifying, distributing, and making available the Artifact; in connection with a Model as Artifact, Use also includes creating content, fine-tuning, updating, running, training, evaluating and re-parametrizing such Model.
+(o) “You” (or “Your”) means an individual or legal entity receiving and exercising permissions granted by this License or making use of the Artifact for the Permitted Purpose and in any permitted field of  use, including usage of the Artifact in an end-use application.
+
+II. INTELLECTUAL PROPERTY RIGHTS
+1. Both copyright and patent grants may apply to the Artifacts. The Artifacts are subject to additional terms as described in Section III below, which govern the Use of the Artifacts in the event that Section II is held unenforceable or inapplicable.
+2. Grant of Copyright License. Conditioned upon compliance with Section III below and subject to the terms and conditions of this License, each Contributor hereby grants to You, only in connection with the Permitted Purpose, a worldwide, non-exclusive, royalty-free copyright license to reproduce, publicly display, publicly perform, distribute, and make derivatives of the Artifacts.
+3. Grant of Patent License. Conditioned upon compliance with Section III below and subject to the terms and conditions of this License, and only where and as applicable, each Contributor hereby grants to You, only in connection with the Permitted Purpose, a worldwide, non-exclusive, royalty-free, irrevocable (except as stated in this paragraph) patent license to make, have made, use, sell, offer to sell, import, and otherwise transfer the Artifacts where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contributions alone or by combination of their Contributions with the Artifact to which such Contribution was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that an Artifact or Contribution constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License in connection with the Artifact shall terminate as of the date such litigation is asserted or filed.
+4. Licensor and Contributor each have the right to grant the licenses above.
+5. The Data is not licensed under this License.
+
+III. CONDITIONS OF USAGE, DISTRIBUTION AND REDISTRIBUTION
+1. Use-based restrictions. The restrictions set forth in Attachment A are mandatory Use-based restrictions. Therefore You may not Use any Artifact in violation of such restrictions. You may Use Artifacts only subject to this License. You shall require all of Your users who Use the Artifacts or Derivatives to comply with the terms of this paragraph and only Use the Artifacts and Derivatives for the Permitted Purpose.
+2. The Output You Generate. Except as set forth herein, Licensor claims no rights in the Output generated by You or Your users. You are accountable for the Output You generate and its subsequent uses. No use of the Output may contravene any provision as stated in this License.
+3. Distribution and Redistribution. You may host for Third Party remote access purposes (e.g. software-as-a-service), reproduce, distribute, and make available Artifacts and Derivatives in any medium, with or without modifications, provided that You meet the following conditions:
+(a) Use-based restrictions in Paragraph III.1 MUST be included as a condition precedent to effect any type of legal agreement (e.g., a license) governing the Use of the Artifacts and Derivatives, and You shall give such notice to any subsequent Third Party recipients.
+(b) You shall give any Third Party recipients of any Artifacts or Derivatives a copy of this License.
+(c) You shall cause any modified files to carry prominent notices stating that You changed the files.
+(d) You shall retain all copyright, patent, trademark, and attribution notices excluding those notices that do not pertain to any part of the Artifacts or Derivatives.
+(e) You and any Third Party recipients of any Artifacts or Derivatives shall adhere to the Permitted Purpose.
+4. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions with respect to Paragraph III.3(a), to govern the Use or Distribution of Your modifications, or for any Derivative, provided that Your Use and Distribution of the Artifacts or their Derivatives otherwise complies with the conditions stated in this License. In other words, the Use-based restrictions referred to in Paragraph III.1 form the minimum set of terms for You to license to Third Parties any Artifacts or Derivatives, but You may add more restrictive terms if You deem it necessary.
+
+IV. OTHER PROVISIONS
+1. Updates and Runtime Restrictions. To the maximum extent permitted by law, Licensor reserves the right to restrict (remotely or otherwise) usage of Artifacts in violation of this License or update Artifacts through electronic means.
+2. Trademarks. Nothing in this License permits You to make use of Licensor’s trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between You and Licensor; and any rights not expressly granted herein are reserved by the Licensor.
+3. DISCLAIMER OF WARRANTY. UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING, LICENSOR PROVIDES THE ARTIFACT (AND EACH CONTRIBUTOR PROVIDES ITS CONTRIBUTIONS) ON AN “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OR CONDITIONS OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING THE ARTIFACT, AND ASSUME ANY RISKS ASSOCIATED WITH YOUR EXERCISE OF PERMISSIONS UNDER THIS LICENSE.
+4. LIMITATION OF LIABILITY. IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW (SUCH AS DELIBERATE AND GROSSLY NEGLIGENT ACTS) OR AGREED TO IN WRITING, SHALL ANY CONTRIBUTOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES OF ANY CHARACTER ARISING AS A RESULT OF THIS LICENSE OR OUT OF THE USE OR INABILITY TO USE THE ARTIFACT (INCLUDING BUT NOT LIMITED TO DAMAGES FOR LOSS OF GOODWILL, WORK STOPPAGE, COMPUTER FAILURE OR MALFUNCTION, OR ANY AND ALL OTHER COMMERCIAL DAMAGES OR LOSSES), EVEN IF SUCH CONTRIBUTOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+5. Severability. If any provision of this License is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein.
+6. Term and Termination. The term of this License will commence upon the earlier of (a) Your acceptance of this License or (b) accessing the Artifact; and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Licensor may terminate this License if You are in breach of any term or condition of this Agreement. Upon termination of this Agreement, You shall delete and cease use of the Artifacts. This paragraph shall survive the termination of this License.
+
+Attachment A – Use Restrictions
+1. Discrimination. You agree not to Use, or allow others to Use, Artifacts or Derivatives
+(a) to discriminate, mock, or promote hatred against individuals or groups, or encourage others to do so directly or indirectly, on the basis of their age; race, perceived race, or ethnicity; national origin; sexual orientation; gender, gender identity, or gender expression; religion or religious affiliation or beliefs; disability status including diseases, bodily conditions, disfigurement, mobility issues, and mental impairment; veteran status; caste; or familial status.
+(b) to exploit any of the vulnerabilities of an individual or specific group of persons based on their age, social, physical, or mental characteristics, in order to materially distort the behavior of a person pertaining to that group in a manner that causes or is likely to cause that person or another person physical or psychological harm.
+(c) to engage in, promote, incite, or facilitate discrimination or other unlawful or harmful conduct in the provision of employment, employment benefits, credit, housing, or other essential goods and services.
+2. Intellectual Property. You agree not to, and not to allow others to Use Artifacts or Derivatives 
+(a) to infringe or attempt to infringe, misappropriate or otherwise violate any intellectual property rights of Licensor or any Third Party; 
+(b) synthesize or modify a natural person’s appearance, voice, or other individual characteristics, unless prior informed consent of said natural person is obtained; or
+(c) to reverse engineer, disassemble, decompile, or otherwise attempt to derive or gain access to Data that was used to create, train, pretrain, or otherwise evaluate such Artifacts or Derivatives.
+3. Legal. You agree not to Use, or allow others to Use, Artifacts or Derivatives
+(a) in any way that violates any applicable national, federal, state, local or international law or regulation;
+(b) to engage in, facilitate, or assist in the planning or development of criminal activities; or
+(c) to generate unlawful content.
+4. Disinformation. You agree not to Use, or allow others to Use, Artifacts or Derivatives
+(a) to create, present or disseminate false or misleading information for economic gain or to intentionally deceive the public, including creating false impersonations of natural persons;
+(b) to defame or harm a person’s reputation, such as by generating, creating, promoting, or spreading defamatory content.
+5. Privacy. You agree not to Use, or allow others to Use, Artifacts or Derivatives
+(a) to engage in, promote, incite, or facilitate the harassment, abuse, threatening or bullying of individuals or groups of individuals; or
+(b) in connection with personal information to infer additional personal       information about a natural person, including but not limited to legally protected characteristics, vulnerabilities or categories; unless informed consent from the data subject to collect said inferred personal information for a stated purpose and defined duration is received.
+6. Health and Safety. You agree not to Use, or allow others to Use, Artifacts or Derivatives
+(a) to provide health or medical advice, medical results interpretation, or make clinical decisions; or
+(b) in connection with any activities that present a risk of death or bodily harm to individuals, including self-harm or harm to others, or in connection with regulated or controlled substances.
+7. Military or Law Enforcement. You agree not to Use, or allow others to Use, Artifacts or Derivatives 
+(a) for purposes of administration of justice, law enforcement, immigration, or asylum processes, such as predicting that a natural person will commit a crime or the likelihood thereof;
+(b) for weaponry or warfare; for building or optimizing weapons; or in service of nuclear proliferation or nuclear weapons technology; or
+(c) military surveillance, including any research or development relating to military surveillance.
+8. General. You agree not to Use, or allow others to Use, Artifacts or Derivatives 
+(a) in any manner that would constitute high risk, restricted, or prohibited of AI under applicable law; or
+(b) to generate or disseminate malware or ransomware or to otherwise harm electronic systems.

cube/README.md

@@ -0,0 +1,166 @@
+# Cube: Generative AI System for 3D
+
+<p align="center">
+  <img src="./resources/teaser.png" width="800" style="margin: 5px;">
+</p>
+
+<div align="center">
+  <a href=https://corp.roblox.com/newsroom/2025/03/introducing-roblox-cube target="_blank"><img src=https://img.shields.io/badge/Roblox-Blog-000000.svg?logo=Roblox height=22px></a>
+  <a href=https://huggingface.co/Roblox/cube3d-0.1 target="_blank"><img src=https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Models-d96902.svg height=22px></a>
+  <a href=https://arxiv.org/abs/2503.15475 target="_blank"><img src=https://img.shields.io/badge/ArXiv-Report-b5212f.svg?logo=arxiv height=22px></a>
+  <a href=https://colab.research.google.com/drive/1ZvTj49pjDCD_crX5WPZNTAoTTzL6-E5t target="_blank"><img src=https://img.shields.io/badge/Google-Open_In_Colab-blue.svg?logo=googlecolab height=22px></a>
+</div>
+
+
+Foundation models trained on vast amounts of data have demonstrated remarkable reasoning and
+generation capabilities in the domains of text, images, audio and video. Our goal is to build
+such a foundation model for 3D intelligence, a model that can support developers in producing all aspects
+of a Roblox experience, from generating 3D objects and scenes to rigging characters for animation to
+producing programmatic scripts describing object behaviors. As we start open-sourcing a family of models 
+towards this vision, we hope to engage others in the research community to address these goals with us.
+
+## Get Started with Cube 3D
+
+<p align="center">
+  <img src="./resources/greyscale_512.gif" width="600" style="margin: 5px;">
+</p>
+
+Cube 3D is our first step towards 3D intelligence, which involves a shape tokenizer and a text-to-shape generation model. We are unlocking the power of generating 3D assets and enhancing creativity for all artists. Our latest version of Cube 3D is now accessible to individuals, creators, researchers and businesses of all sizes so that they can experiment, innovate and scale their ideas responsibly. This release includes model weights and starting code for using our text-to-shape model to create 3D assets.
+
+### Try it out on [Google Colab](https://colab.research.google.com/drive/1ZvTj49pjDCD_crX5WPZNTAoTTzL6-E5t)
+
+### Install Requirements
+
+Clone and install this repo in a virtual environment, via:
+
+```bash
+git clone https://github.com/Roblox/cube.git
+cd cube
+pip install -e .[meshlab]
+```
+
+> **CUDA**: If you are using a Windows machine, you may need to install the [CUDA](https://developer.nvidia.com/cuda-downloads) toolkit as well as `torch` with cuda support via `pip install torch --index-url https://download.pytorch.org/whl/cu124 --force-reinstall`
+
+> **Note**: `[meshlab]` is an optional dependency and can be removed by simply running `pip install -e .` for better compatibility but mesh simplification will be disabled.
+
+### Download Models from Huggingface 🤗
+
+Download the model weights from [hugging face](https://huggingface.co/Roblox/cube3d-v0.1) or use the
+`huggingface-cli`:
+
+```bash
+huggingface-cli download Roblox/cube3d-v0.1 --local-dir ./model_weights
+```
+
+### Inference
+
+#### 1. Shape Generation
+
+To generate 3D models using the downloaded models simply run:
+
+```bash
+python -m cube3d.generate \
+            --gpt-ckpt-path model_weights/shape_gpt.safetensors \
+            --shape-ckpt-path model_weights/shape_tokenizer.safetensors \
+            --fast-inference \
+            --prompt "Broad-winged flying red dragon, elongated, folded legs."
+```
+
+> **Note**: `--fast-inference` is optional and may not be available for all GPU that have limited VRAM. This flag will also not work on MacOS. 
+
+The output will be an `.obj` file saved in the specified `output` directory.
+
+If you want to render a turntable gif of the mesh, you can use the `--render-gif` flag, which will render a turntable gif of the mesh
+and save it as `turntable.gif` in the specified `output` directory. 
+
+We provide several example output objects and their corresponding text prompts in the `examples` folder.
+
+> **Note**: You must have Blender installed and available in your system's PATH to render the turntable GIF. You can download it from [Blender's official website](https://www.blender.org/). Ensure that the Blender executable is accessible from the command line.
+
+> **Note**: If shape decoding is slow, you can try to specify a lower resolution using the `--resolution-base` flag. A lower resolution will create a coarser and lower quality output mesh but faster decoding. Values between 4.0 and 9.0 are recommended.
+
+#### 2. Shape Tokenization and De-tokenization
+
+To tokenize a 3D shape into token indices and reconstruct it back, you can use the following command:
+
+```bash
+python -m cube3d.vq_vae_encode_decode \
+            --shape-ckpt-path model_weights/shape_tokenizer.safetensors \
+            --mesh-path ./outputs/output.obj
+```
+
+This will process the `.obj` file located at `./outputs/output.obj` and prints the tokenized representation as well as exports the mesh reconstructed from the token indices.
+
+### Hardware Requirements
+
+We have tested our model on:
+* Nvidia H100 GPU
+* Nvidia A100 GPU
+* Nvidia Geforce 3080
+* Apple Silicon M2-4 Chips.
+
+We recommend using a GPU with at least 24GB of VRAM available when using `--fast-inference` (or `EngineFast`) and 16GB otherwise. 
+
+### Code Usage
+
+We have designed a minimalist API that allows the use this repo as a Python library:
+
+```python
+import torch
+import trimesh
+from cube3d.inference.engine import Engine, EngineFast
+
+# load ckpt
+config_path = "cube3d/configs/open_model.yaml"
+gpt_ckpt_path = "model_weights/shape_gpt.safetensors"
+shape_ckpt_path = "model_weights/shape_tokenizer.safetensors"
+engine_fast = EngineFast( # only supported on CUDA devices, replace with Engine otherwise
+    config_path, 
+    gpt_ckpt_path, 
+    shape_ckpt_path, 
+    device=torch.device("cuda"),
+)
+
+# inference
+input_prompt = "A pair of noise-canceling headphones"
+# NOTE: Reduce `resolution_base` for faster inference and lower VRAM usage
+# The `top_k` parameter controls randomness between inferences:
+#   - A value of 1 yields deterministic results.
+#   - Higher values introduce more randomness.
+mesh_v_f = engine_fast.t2s([input_prompt], use_kv_cache=True, resolution_base=8.0, top_k=5)
+
+# save output
+vertices, faces = mesh_v_f[0][0], mesh_v_f[0][1]
+_ = trimesh.Trimesh(vertices=vertices, faces=faces).export("output.obj")
+```
+
+## Coming Soon
+
+### Controlling shape generation with bounding box conditioning
+<div align="center">
+  <img src="./resources/truck_black_text_512.gif" width="300" height="300" style="margin: 5px;">
+  <img src="./resources/couch_black_text_512.gif" width="300" height="300" style="margin: 5px;">
+</div>
+
+### Scene Generation
+
+https://github.com/user-attachments/assets/987c459a-5708-41a5-9b92-89068a70a239
+
+https://github.com/user-attachments/assets/ab501a86-b0cb-4c73-827e-988b2120d4c0
+
+## Citation
+If you find this work helpful, please consider citing our technical report:
+
+```bibtex
+@article{roblox2025cube,
+    title   = {Cube: A Roblox View of 3D Intelligence},
+    author  = {Roblox, Foundation AI Team},
+    journal = {arXiv preprint arXiv:2503.15475},
+    year    = {2025}
+}
+```
+
+## Acknowledgements
+
+We would like to thank the contributors of [TRELLIS](https://github.com/microsoft/TRELLIS), [CraftsMan3D](https://github.com/wyysf-98/CraftsMan3D), [threestudio](https://github.com/threestudio-project/threestudio), [Hunyuan3D-2](https://github.com/Tencent/Hunyuan3D-2), [minGPT](https://github.com/karpathy/minGPT), [dinov2](https://github.com/facebookresearch/dinov2), [OptVQ](https://github.com/zbr17/OptVQ), [1d-tokenizer](https://github.com/bytedance/1d-tokenizer)
+repositories, for their open source contributions.

cube/SECURITY.md

@@ -0,0 +1,9 @@
+# Security Policy
+
+## Reporting a Vulnerability
+
+If you discover a security vulnerability in this repository, we appreciate your help in ensuring that the issue is addressed quickly. 
+
+Report any vulnerabilities found to our bug bounty program on HackerOne: https://hackerone.com/roblox
+
+Please **do not create a public issue in this repo**. 

cube/cube3d/configs/open_model.yaml

@@ -0,0 +1,32 @@
+gpt_model:
+  n_layer: 23
+  n_single_layer: 1
+  rope_theta: 10000
+  n_head: 12
+  n_embd: 1536
+  bias: true
+  eps: 1.e-6
+  shape_model_vocab_size: 16384
+  text_model_embed_dim: 768
+  use_pooled_text_embed: False
+  shape_model_embed_dim: 32
+  encoder_with_cls_token: true
+
+shape_model:
+  encoder_with_cls_token: true
+  num_encoder_latents: 512 
+  num_decoder_latents: 0 
+  embed_dim: 32
+  width: 768
+  num_heads: 12
+  out_dim: 1
+  eps: 1.e-6 
+  num_freqs: 128 
+  point_feats: 3 
+  embed_point_feats: false
+  num_encoder_layers: 13
+  encoder_cross_attention_levels: [0, 2, 4, 8] 
+  num_decoder_layers: 24
+  num_codes: 16384
+
+text_model_pretrained_model_name_or_path: "openai/clip-vit-large-patch14"

cube/cube3d/generate.py

@@ -0,0 +1,144 @@
+import argparse
+import os
+
+import torch
+import trimesh
+
+from cube3d.inference.engine import Engine, EngineFast
+from cube3d.mesh_utils.postprocessing import (
+    PYMESHLAB_AVAILABLE,
+    create_pymeshset,
+    postprocess_mesh,
+    save_mesh,
+)
+from cube3d.renderer import renderer
+
+def generate_mesh(
+    engine,
+    prompt,
+    output_dir,
+    output_name,
+    resolution_base=8.0,
+    disable_postprocess=False,
+    top_k: int = 1,
+):
+    mesh_v_f = engine.t2s(
+        [prompt],
+        use_kv_cache=True,
+        resolution_base=resolution_base,
+        top_k=top_k,
+    )
+    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
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="cube shape generation script")
+    parser.add_argument(
+        "--config-path",
+        type=str,
+        default="cube3d/configs/open_model.yaml",
+        help="Path to the configuration YAML 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(
+        "--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-k",
+        type=int,
+        default=1,
+        help="Top k filtering, 0 means no filtering, by default 1, which is determistic.",
+    )
+    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 = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    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, device=device
+        )
+    
+    # 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_k,
+    )
+    if args.render_gif:
+        gif_path = renderer.render_turntable(obj_path, args.output_dir)
+        print(f"Rendered turntable gif for {args.prompt} at `{gif_path}`")
+    print(f"Generated mesh for {args.prompt} at `{obj_path}`")

cube/cube3d/inference/engine.py

@@ -0,0 +1,499 @@
+import torch
+from tqdm import tqdm
+from transformers import CLIPTextModelWithProjection, CLIPTokenizerFast
+
+from cube3d.inference.logits_postprocesses import process_logits
+from cube3d.inference.utils import load_config, load_model_weights, parse_structured
+from cube3d.model.autoencoder.one_d_autoencoder import OneDAutoEncoder
+from cube3d.model.gpt.dual_stream_roformer import DualStreamRoformer
+from cube3d.model.transformers.cache import Cache
+
+
+class Engine:
+    def __init__(
+        self,
+        config_path: str,
+        gpt_ckpt_path: str,
+        shape_ckpt_path: str,
+        device: torch.device,
+    ):
+        """
+        Initializes the inference engine with the given configuration and checkpoint paths.
+        Args:
+            config_path (str): Path to the configuration file.
+            gpt_ckpt_path (str): Path to the GPT model checkpoint file.
+            shape_ckpt_path (str): Path to the shape model checkpoint file.
+            device (torch.device): The device to run the models on (e.g., 'cpu' or 'cuda').
+        Attributes:
+            cfg (dict): Loaded configuration from the config file.
+            device (torch.device): The device to run the models on.
+            gpt_model (DualStreamRoformer): The GPT model initialized and loaded with weights.
+            shape_model (OneDAutoEncoder): The shape model initialized and loaded with weights.
+            text_model (CLIPTextModelWithProjection): The text model initialized from a pretrained model.
+            text_tokenizer (CLIPTokenizerFast): The tokenizer for the text model.
+            max_new_tokens (int): Maximum number of new tokens for the shape model.
+            min_id (int): Minimum ID for the shape model codes.
+            max_id (int): Maximum ID for the shape model codes.
+        """
+
+        self.cfg = load_config(config_path)
+        self.device = device
+
+        self.gpt_model = DualStreamRoformer(
+            parse_structured(DualStreamRoformer.Config, self.cfg.gpt_model)
+        )
+        load_model_weights(
+            self.gpt_model,
+            gpt_ckpt_path,
+        )
+        self.gpt_model = self.gpt_model.eval().to(self.device)
+
+        self.shape_model = OneDAutoEncoder(
+            parse_structured(OneDAutoEncoder.Config, self.cfg.shape_model)
+        )
+        load_model_weights(
+            self.shape_model,
+            shape_ckpt_path,
+        )
+        self.shape_model = self.shape_model.eval().to(self.device)
+
+        # copy vq codebook to gpt
+        with torch.no_grad():
+            codebook = self.shape_model.bottleneck.block.get_codebook()
+            codebook = self.gpt_model.shape_proj(codebook).detach()
+        self.gpt_model.transformer.wte.weight.data[: codebook.shape[0]] = codebook
+
+        self.text_model = CLIPTextModelWithProjection.from_pretrained(
+            self.cfg.text_model_pretrained_model_name_or_path,
+            force_download=False,
+            device_map=self.device,
+        ).eval()
+        self.text_tokenizer = CLIPTokenizerFast.from_pretrained(
+            self.cfg.text_model_pretrained_model_name_or_path
+        )
+
+        self.max_new_tokens = self.shape_model.cfg.num_encoder_latents
+        self.min_id = 0
+        self.max_id = self.shape_model.cfg.num_codes
+
+    @torch.inference_mode()
+    def prepare_inputs(self, prompts: list[str], guidance_scale: float):
+        """
+        Prepares the input embeddings for the model based on the provided prompts and guidance scale.
+        Args:
+            prompts (list[str]): A list of prompt strings to be encoded.
+            guidance_scale (float): A scaling factor for guidance. If greater than 0.0, additional processing is applied.
+        Returns:
+            tuple: A tuple containing:
+                - embed (torch.Tensor): The encoded input embeddings.
+                - cond (torch.Tensor): The condition embeddings, which may include unconditional embeddings if guidance_scale is greater than 0.0.
+        """
+
+        prompt_embeds = self.run_clip(prompts)
+
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            embed = self.encode_input(prompt_embeds, self.gpt_model.shape_bos_id)
+
+        cond = prompt_embeds
+        if guidance_scale > 0.0:
+            embed = torch.cat([embed, embed], dim=0)
+            uncond_embeds = self.run_clip([""] * len(prompts))
+            cond = torch.cat([prompt_embeds, uncond_embeds], dim=0)
+
+        return embed, cond
+
+    @torch.inference_mode()
+    def run_clip(self, text_inputs):
+        """
+        Processes the given text inputs using a text tokenizer and a text model, and returns the encoded text embeddings.
+        Args:
+            text_inputs (str or List[str]): The input text or list of texts to be processed.
+        Returns:
+            torch.Tensor: The encoded text embeddings.
+        """
+
+        text_inputs = self.text_tokenizer(
+            text_inputs,
+            max_length=self.text_tokenizer.model_max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+        with torch.no_grad():
+            text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}
+            # use full precision for text encoder
+            with torch.autocast(device_type=self.device.type, enabled=False):
+                encoded = self.text_model(**text_inputs)
+            if self.gpt_model.cfg.use_pooled_text_embed:
+                embed = encoded.text_embeds.unsqueeze(1)  # [bs, 1, 512]
+            else:
+                embed = encoded.last_hidden_state  # [bs, 77, 512]
+        embed = self.gpt_model.encode_text(embed)
+
+        return embed
+
+    @torch.inference_mode()
+    def encode_input(self, inputs: torch.Tensor, bos: int):
+        """
+        Encodes the beginning of sequence (BOS) token for the given input tensor.
+        Args:
+            inputs (torch.Tensor): The input tensor containing sequences.
+            bos (int): The beginning of sequence token ID.
+        Returns:
+            torch.Tensor: The encoded BOS token embeddings.
+        """
+
+        b = inputs.shape[0]
+        bos_embed = self.gpt_model.encode_token(
+            torch.full(
+                (b, 1),
+                fill_value=bos,
+                dtype=torch.long,
+                device=self.device,
+            )
+        )
+        return bos_embed
+
+    @torch.inference_mode()
+    def run_gpt(
+        self,
+        prompts: list[str],
+        use_kv_cache: bool,
+        guidance_scale: float = 3.0,
+        top_k: int = 1,
+    ):
+        """
+        Generates text using a GPT model based on the provided prompts.
+        Args:
+            prompts (list[str]): A list of input prompts to generate text from.
+            use_kv_cache (bool): Whether to use key-value caching for faster generation.
+            guidance_scale (float, optional): The scale for guidance during generation. Default is 3.0.
+            top_k : (int, optional): Top k filtering, 0 means no filtering, by default 1.
+        Returns:
+            torch.Tensor: A tensor containing the generated token IDs.
+        """
+        embed, cond = self.prepare_inputs(prompts, guidance_scale)
+
+        output_ids = []
+
+        batch_size, input_seq_len, dim = embed.shape
+        max_seq_len = input_seq_len + self.max_new_tokens
+        embed_buffer = torch.zeros(
+            (batch_size, max_seq_len, dim), dtype=embed.dtype, device=embed.device
+        )
+        embed_buffer[:, :input_seq_len, :].copy_(embed)
+        cond_len = cond.shape[1]
+        kv_cache = None
+        if use_kv_cache:
+            kv_cache = self.gpt_model.init_kv_cache(
+                batch_size,
+                cond_len,
+                self.max_new_tokens + 1,  # +1 for the BOS token
+                torch.bfloat16,
+                embed.device,
+            )
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            for i in tqdm(range(self.max_new_tokens), desc=f"generating"):
+                curr_pos_id = torch.tensor([i], dtype=torch.long, device=embed.device)
+                logits = self.gpt_model(
+                    embed_buffer,
+                    cond,
+                    kv_cache=kv_cache,
+                    curr_pos_id=curr_pos_id if use_kv_cache else None,
+                    decode=(i > 0) if use_kv_cache else False,
+                )
+                if use_kv_cache:
+                    logits = logits[:, 0, ...]
+                else:
+                    logits = logits[:, i, ...]
+
+                logits = logits[..., self.min_id : self.max_id]
+
+                if guidance_scale > 0.0:
+                    logits, uncond_logits = logits.float().chunk(2, dim=0)
+                    gamma = (
+                        guidance_scale * (self.max_new_tokens - i) / self.max_new_tokens
+                    )
+                    logits = (1 + gamma) * logits - gamma * uncond_logits
+                probs = process_logits(
+                    logits,
+                    top_k=top_k,
+                )
+                next_id = torch.multinomial(probs, num_samples=1, replacement=True)
+                output_ids.append(next_id)
+                next_embed = self.gpt_model.encode_token(next_id)
+                if guidance_scale > 0.0:
+                    next_embed = torch.cat([next_embed, next_embed], dim=0)
+                embed_buffer[:, i + input_seq_len, :].copy_(next_embed.squeeze(1))
+
+        return torch.cat(output_ids, dim=1)
+
+    @torch.inference_mode()
+    def run_shape_decode(
+        self,
+        output_ids: torch.Tensor,
+        resolution_base: float = 8.0,
+        chunk_size: int = 100_000,
+    ):
+        """
+        Decodes the shape from the given output IDs and extracts the geometry.
+        Args:
+            output_ids (torch.Tensor): The tensor containing the output IDs.
+            resolution_base (float, optional): The base resolution for geometry extraction. Defaults to 8.43.
+            chunk_size (int, optional): The chunk size for processing. Defaults to 100,000.
+        Returns:
+            tuple: A tuple containing the vertices and faces of the mesh.
+        """
+        shape_ids = (
+            output_ids[:, : self.shape_model.cfg.num_encoder_latents, ...]
+            .clamp_(0, self.shape_model.cfg.num_codes - 1)
+            .view(-1, self.shape_model.cfg.num_encoder_latents)
+        )
+        latents = self.shape_model.decode_indices(shape_ids)
+        mesh_v_f, _ = self.shape_model.extract_geometry(
+            latents,
+            resolution_base=resolution_base,
+            chunk_size=chunk_size,
+            use_warp=True,
+        )
+        return mesh_v_f
+
+    @torch.inference_mode()
+    def t2s(
+        self,
+        prompts: list[str],
+        use_kv_cache: bool,
+        guidance_scale: float = 3.0,
+        resolution_base: float = 8.0,
+        chunk_size: int = 100_000,
+        top_k: int = 5,
+    ):
+        """
+        Generates a 3D mesh from text prompts using a GPT model and shape decoder.
+        Args:
+            prompts (list[str]): A list of text prompts to guide the generation.
+            use_kv_cache (bool): Whether to use key-value caching for the GPT model.
+            guidance_scale (float, optional): The scale of guidance for the GPT model. Default is 3.0.
+            resolution_base (float, optional): The base resolution for the shape decoder. Default is 8.0.
+            chunk_size (int, optional): The chunk size for processing the shape decoding. Default is 100,000.
+        Returns:
+            mesh_v_f: The generated 3D mesh vertices and faces.
+        """
+        output_ids = self.run_gpt(prompts, use_kv_cache, guidance_scale, top_k)
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            mesh_v_f = self.run_shape_decode(output_ids, resolution_base, chunk_size)
+        return mesh_v_f
+
+
+class EngineFast(Engine):
+    def __init__(
+        self,
+        config_path: str,
+        gpt_ckpt_path: str,
+        shape_ckpt_path: str,
+        device: torch.device,
+    ):
+        """
+        Initializes the inference engine with the given configuration and checkpoint paths.
+        Args:
+            config_path (str): Path to the configuration file.
+            gpt_ckpt_path (str): Path to the GPT checkpoint file.
+            shape_ckpt_path (str): Path to the shape checkpoint file.
+            device (torch.device): The device to run the inference on (e.g., CPU or CUDA).
+        """
+
+        super().__init__(config_path, gpt_ckpt_path, shape_ckpt_path, device)
+
+        # CUDA Graph params
+        self.graph = torch.cuda.CUDAGraph()
+        self.embed_buffer = torch.Tensor()
+        self.cond_buffer = torch.Tensor()
+        self.logits_buffer = torch.Tensor()
+        self.curr_pos_id = torch.tensor([0], dtype=torch.long, device=self.device)
+        self.kv_cache: list[Cache] = []
+
+        self._warmup_and_capture_graph()
+
+    def _warmup_and_capture_graph(self):
+        """
+        Warms up the model by running a series of forward passes and captures the CUDA graph for efficient execution.
+        This method performs the following steps:
+        1. Prepares the input embeddings and conditions using a warmup prompt.
+        2. Initializes buffers for embeddings and conditions.
+        3. Initializes the key-value cache for the GPT model.
+        4. Runs a series of warmup passes to prefill the model and generate logits.
+        5. Captures the CUDA graph for the model's forward pass to optimize future executions.
+        """
+
+        warmup_prompt = "A cube"
+        embed, cond = self.prepare_inputs([warmup_prompt], guidance_scale=3.0)
+
+        batch_size, input_seq_len, dim = embed.shape
+        max_seq_len = input_seq_len + self.max_new_tokens
+        self.embed_buffer = torch.zeros(
+            (batch_size, max_seq_len, dim), dtype=embed.dtype, device=self.device
+        )
+        self.embed_buffer[:, :input_seq_len, :].copy_(embed)
+
+        self.cond_buffer = torch.empty_like(cond)
+        self.cond_buffer.copy_(cond)
+        cond_len = self.cond_buffer.shape[1]
+
+        # Initialize kv_cache for the first time
+        self.kv_cache = self.gpt_model.init_kv_cache(
+            batch_size,
+            cond_len,
+            self.max_new_tokens + 1,  # +1 for the BOS token
+            torch.bfloat16,
+            self.device,
+        )
+
+        num_warmup_passes = 10
+
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            self._set_curr_pos_id(0)
+            _ = self._prefill_and_return_logits()
+
+            for x in range(1, num_warmup_passes):
+                self._set_curr_pos_id(x)
+                self.logits_buffer = self.gpt_model(
+                    embed=self.embed_buffer,
+                    cond=self.cond_buffer,
+                    kv_cache=self.kv_cache,
+                    curr_pos_id=self.curr_pos_id,
+                    decode=True,
+                )
+
+        side_stream = torch.cuda.Stream(device=self.device)
+        with torch.cuda.graph(self.graph, stream=side_stream):
+            with torch.autocast(self.device.type, dtype=torch.bfloat16):
+                self.logits_buffer = self.gpt_model(
+                    embed=self.embed_buffer,
+                    cond=self.cond_buffer,
+                    kv_cache=self.kv_cache,
+                    curr_pos_id=self.curr_pos_id,
+                    decode=True,
+                )
+
+    def _reset_kv_cache(self):
+        """
+        Resets the key-value cache by setting all key and value states to zero.
+        This method iterates through each cache in the `kv_cache` attribute and
+        calls the `zero_()` method on both `key_states` and `value_states` to
+        reset them to their initial state.
+        """
+
+        for cache in self.kv_cache:
+            cache.key_states.zero_()
+            cache.value_states.zero_()
+
+    def _prefill_and_return_logits(self) -> torch.Tensor:
+        """
+        Prefills the model's key-value cache and returns the logits.
+        This method resets the key-value cache and then performs a forward pass
+        through the GPT model in eager mode to prefill the logits.
+        Returns:
+            torch.Tensor: The prefilled logits tensor with the first dimension removed.
+        """
+
+        self._reset_kv_cache()
+
+        # Prefill is always eager
+        prefill_logits = self.gpt_model(
+            embed=self.embed_buffer,
+            cond=self.cond_buffer,
+            kv_cache=self.kv_cache,
+            curr_pos_id=self.curr_pos_id,
+            decode=False,
+        )
+
+        return prefill_logits[:, 0, ...]
+
+    def _set_curr_pos_id(self, pos: int):
+        """
+        Set the current position ID.
+        This method updates the `curr_pos_id` attribute with the given position.
+        Args:
+            pos (int): The position ID to set.
+        """
+
+        self.curr_pos_id.copy_(
+            torch.tensor([pos], dtype=torch.long, device=self.device)
+        )
+
+    def run_gpt(
+        self, 
+        prompts: list[str], 
+        use_kv_cache: bool, 
+        guidance_scale: float = 3.0,
+        top_k: int = 1,
+    ):
+        """
+        Runs the GPT model to generate text based on the provided prompts.
+        Args:
+            prompts (list[str]): A list of input prompts for the GPT model. Only a single prompt is supported.
+            use_kv_cache (bool): Flag indicating whether to use key-value caching. (Currently not used)
+            guidance_scale (float, optional): The scale factor for guidance. Default is 3.0.
+        Returns:
+            torch.Tensor: A tensor containing the generated output token IDs.
+        Raises:
+            AssertionError: If the batch size is greater than 1.
+        """
+
+        embed, cond = self.prepare_inputs(prompts, guidance_scale)
+        assert len(prompts) == 1, "batch size > 1 not support for EngineFast"
+
+        batch_size, input_seq_len, _ = embed.shape
+        self.embed_buffer.zero_()
+        self.embed_buffer[:, :input_seq_len, :].copy_(embed)
+
+        assert self.cond_buffer.shape == cond.shape
+        self.cond_buffer.copy_(cond)
+
+        output_ids = torch.zeros(
+            (batch_size // 2, self.max_new_tokens), dtype=torch.int, device=self.device
+        )
+
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            self._set_curr_pos_id(0)
+
+            logits = self._prefill_and_return_logits()
+
+            logits = logits[..., self.min_id : self.max_id]
+            if guidance_scale > 0.0:
+                logits, uncond_logits = logits.float().chunk(2, dim=0)
+                gamma = guidance_scale
+                logits = (1 + gamma) * logits - gamma * uncond_logits
+
+            probs = process_logits(logits, top_k=top_k)
+            next_id = torch.multinomial(probs, num_samples=1, replacement=True)
+
+            output_ids[:, 0] = next_id.squeeze()
+            next_embed = self.gpt_model.encode_token(next_id)
+            next_embed = next_embed.repeat(2, 1, 1)
+            self.embed_buffer[:, input_seq_len, :].copy_(next_embed.squeeze(1))
+
+            for i in tqdm(
+                range(1, self.max_new_tokens), desc=f"generating"
+            ):
+                self._set_curr_pos_id(i)
+                self.graph.replay()
+
+                logits = self.logits_buffer[:, 0, ...]
+
+                logits = logits[..., self.min_id : self.max_id]
+                if guidance_scale > 0.0:
+                    logits, uncond_logits = logits.float().chunk(2, dim=0)
+                    gamma = (
+                        guidance_scale * (self.max_new_tokens - i) / self.max_new_tokens
+                    )
+                    logits = (1 + gamma) * logits - gamma * uncond_logits
+                probs = process_logits(logits, top_k=top_k)
+                next_id = torch.multinomial(probs, num_samples=1, replacement=True)
+
+                output_ids[:, i] = next_id.squeeze()
+                next_embed = self.gpt_model.encode_token(next_id)
+                next_embed = next_embed.repeat(2, 1, 1)
+                self.embed_buffer[:, i + input_seq_len, :].copy_(next_embed.squeeze(1))
+
+        return output_ids

cube/cube3d/inference/logits_postprocesses.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn.functional as F
+
+
+def top_k_filtering(logits, top_k: int = 1):
+    """
+    Filter a distribution of logits using top-k and/or top-p (nucleus) filtering.
+    The input logits tensor is modified in-place.
+
+    Args:
+        logits: A tensor of logits to be filtered. Expected shape is [..., vocab_size].
+        top_k: If > 0, only keep the top k tokens with highest probability.
+        top_p: If < 1.0, only keep tokens whose cumulative probability is below this threshold.
+
+    Returns:
+        A tensor of logits where values outside the top-k/top-p threshold are set to -∞.
+    """
+    if top_k > 0:
+        idx_to_remove = logits < logits.topk(top_k, largest=True, sorted=False, dim=-1)[
+            0
+        ].amin(dim=-1, keepdim=True)
+        logits.masked_fill_(idx_to_remove, -torch.inf)
+
+    return logits
+
+
+def process_logits(
+        logits,
+        top_k: int = 1,
+    ):
+    """
+    Process logits by optionally applying top-k filtering.
+    The final probabilities are returned after applying softmax on the filtered logits.
+
+    Args:
+        logits: A tensor of logits to process. Expected shape is [..., vocab_size].
+        top_k: If > 0, only keep the top k tokens with highest probability.
+
+    Returns:
+        A tensor of probabilities after filtering, with the same shape as the input logits.
+    """
+    logits = top_k_filtering(logits, top_k=top_k)
+    probs = F.softmax(logits, dim=-1)
+    return probs

cube/cube3d/inference/utils.py

@@ -0,0 +1,56 @@
+import logging
+from typing import Any, Optional
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+from safetensors.torch import load_model
+
+
+def load_config(cfg_path: str) -> Any:
+    """
+    Load and resolve a configuration file.
+    Args:
+        cfg_path (str): The path to the configuration file.
+    Returns:
+        Any: The loaded and resolved configuration object.
+    Raises:
+        AssertionError: If the loaded configuration is not an instance of DictConfig.
+    """
+
+    cfg = OmegaConf.load(cfg_path)
+    OmegaConf.resolve(cfg)
+    assert isinstance(cfg, DictConfig)
+    return cfg
+
+
+def parse_structured(cfg_type: Any, cfg: DictConfig) -> Any:
+    """
+    Parses a configuration dictionary into a structured configuration object.
+    Args:
+        cfg_type (Any): The type of the structured configuration object.
+        cfg (DictConfig): The configuration dictionary to be parsed.
+    Returns:
+        Any: The structured configuration object created from the dictionary.
+    """
+
+    scfg = OmegaConf.structured(cfg_type(**cfg))
+    return scfg
+
+
+def load_model_weights(model: torch.nn.Module, ckpt_path: str) -> None:
+    """
+    Load a safetensors checkpoint into a PyTorch model.
+    The model is updated in place.
+
+    Args:
+        model: PyTorch model to load weights into
+        ckpt_path: Path to the safetensors checkpoint file
+
+    Returns:
+        None
+    """
+    assert ckpt_path.endswith(".safetensors"), (
+        f"Checkpoint path '{ckpt_path}' is not a safetensors file"
+    )
+
+    load_model(model, ckpt_path)

cube/cube3d/mesh_utils/postprocessing.py

@@ -0,0 +1,84 @@
+import logging
+
+import numpy as np
+
+try:
+    import pymeshlab
+
+    PYMESHLAB_AVAILABLE = True
+except ImportError:
+    logging.warning(
+        "pymeshlab is not installed or could not be loaded. Please install it with `pip install pymeshlab`."
+    )
+    PYMESHLAB_AVAILABLE = False
+    from typing import Any
+
+    # Create stub class for typing
+    class pymeshlab:
+        MeshSet = Any
+        Mesh = Any
+
+
+def create_pymeshset(vertices: np.ndarray, faces: np.ndarray):
+    """
+    Creates a MeshLab MeshSet given a list of vertices and faces.
+    """
+    assert PYMESHLAB_AVAILABLE, "pymeshlab is not installed or could not be loaded."
+    # Initialize MeshSet and create pymeshlab.Mesh
+    mesh_set = pymeshlab.MeshSet()
+    input_mesh = pymeshlab.Mesh(vertex_matrix=vertices, face_matrix=faces)
+    mesh_set.add_mesh(input_mesh, "input_mesh")
+    logging.info("Mesh successfully added to pymeshlab MeshSet.")
+    return mesh_set
+
+
+def cleanup(ms: pymeshlab.MeshSet):
+    """
+    General cleanup for a given Mesh. Removes degenerate elements from the
+    geometry.
+    """
+    ms.meshing_remove_null_faces()
+    ms.meshing_remove_folded_faces()
+    ms.meshing_remove_duplicate_vertices()
+    ms.meshing_remove_duplicate_faces()
+    ms.meshing_remove_t_vertices()
+    ms.meshing_remove_unreferenced_vertices()
+
+
+def remove_floaters(ms: pymeshlab.MeshSet, threshold: float = 0.005):
+    """
+    Remove any floating artifacts that exist from our mesh generation.
+    """
+    assert PYMESHLAB_AVAILABLE, "pymeshlab is not installed or could not be loaded."
+    ms.meshing_remove_connected_component_by_diameter(
+        mincomponentdiag=pymeshlab.PercentageValue(15), removeunref=True
+    )
+
+
+def simplify_mesh(ms: pymeshlab.MeshSet, target_face_num: int):
+    """
+    Simplify the mesh to the target number of faces.
+    """
+    ms.meshing_decimation_quadric_edge_collapse(
+        targetfacenum=target_face_num,
+        qualitythr=0.4,
+        preservenormal=True,
+        autoclean=True,
+    )
+
+
+def save_mesh(ms: pymeshlab.MeshSet, output_path: str):
+    """
+    Save the mesh to a file.
+    """
+    ms.save_current_mesh(output_path)
+    logging.info(f"Mesh saved to {output_path}.")
+
+
+def postprocess_mesh(ms: pymeshlab.MeshSet, target_face_num: int, output_path: str):
+    """
+    Postprocess the mesh to the target number of faces.
+    """
+    cleanup(ms)
+    remove_floaters(ms)
+    simplify_mesh(ms, target_face_num)

cube/cube3d/model/autoencoder/embedder.py

@@ -0,0 +1,52 @@
+import math
+
+import torch
+import torch.nn as nn
+
+
+class PhaseModulatedFourierEmbedder(torch.nn.Module):
+    def __init__(
+        self,
+        num_freqs: int,
+        input_dim: int = 3,
+    ):
+        """
+        Initializes the PhaseModulatedFourierEmbedder class.
+        Args:
+            num_freqs (int): The number of frequencies to be used.
+            input_dim (int, optional): The dimension of the input. Defaults to 3.
+        Attributes:
+            weight (torch.nn.Parameter): The weight parameter initialized with random values.
+            carrier (torch.Tensor): The carrier frequencies calculated based on the Nyquist-Shannon sampling theorem.
+            out_dim (int): The output dimension calculated based on the input dimension and number of frequencies.
+        """
+
+        super().__init__()
+
+        self.weight = nn.Parameter(
+            torch.randn(input_dim, num_freqs) * math.sqrt(0.5 * num_freqs)
+        )
+
+        # NOTE this is the highest frequency we can get (2 for peaks, 2 for zeros, and 4 for interpolation points), see also https://en.wikipedia.org/wiki/Nyquist%E2%80%93Shannon_sampling_theorem
+        carrier = (num_freqs / 8) ** torch.linspace(1, 0, num_freqs)
+        carrier = (carrier + torch.linspace(0, 1, num_freqs)) * 2 * torch.pi
+        self.register_buffer("carrier", carrier, persistent=False)
+
+        self.out_dim = input_dim * (num_freqs * 2 + 1)
+
+    def forward(self, x):
+        """
+        Perform the forward pass of the embedder model.
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, ..., input_dim).
+        Returns:
+            torch.Tensor: Output tensor of shape (batch_size, ..., output_dim) where
+                          output_dim = input_dim + 2 * input_dim.
+        """
+
+        m = x.float().unsqueeze(-1)
+        fm = (m * self.weight).view(*x.shape[:-1], -1)
+        pm = (m * 0.5 * torch.pi + self.carrier).view(*x.shape[:-1], -1)
+        embedding = torch.cat([x, fm.cos() + pm.cos(), fm.sin() + pm.sin()], dim=-1)
+
+        return embedding

cube/cube3d/model/autoencoder/grid.py

@@ -0,0 +1,84 @@
+from typing import Literal, Union
+
+import numpy as np
+import torch
+import warp as wp
+
+
+def generate_dense_grid_points(
+    bbox_min: np.ndarray,
+    bbox_max: np.ndarray,
+    resolution_base: float,
+    indexing: Literal["xy", "ij"] = "ij",
+) -> tuple[np.ndarray, list[int], np.ndarray]:
+    """
+    Generate a dense grid of points within a bounding box.
+
+    Parameters:
+        bbox_min (np.ndarray): The minimum coordinates of the bounding box (3D).
+        bbox_max (np.ndarray): The maximum coordinates of the bounding box (3D).
+        resolution_base (float): The base resolution for the grid. The number of cells along each axis will be 2^resolution_base.
+        indexing (Literal["xy", "ij"], optional): The indexing convention for the grid. "xy" for Cartesian indexing, "ij" for matrix indexing. Default is "ij".
+    Returns:
+        tuple: A tuple containing:
+            - xyz (np.ndarray): A 2D array of shape (N, 3) where N is the total number of grid points. Each row represents the (x, y, z) coordinates of a grid point.
+            - grid_size (list): A list of three integers representing the number of grid points along each axis.
+            - length (np.ndarray): The length of the bounding box along each axis.
+    """
+    length = bbox_max - bbox_min
+    num_cells = np.exp2(resolution_base)
+    x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32)
+    y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32)
+    z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32)
+    [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing)
+    xyz = np.stack((xs, ys, zs), axis=-1)
+    xyz = xyz.reshape(-1, 3)
+    grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1]
+
+    return xyz, grid_size, length
+
+
+def marching_cubes_with_warp(
+    grid_logits: torch.Tensor,
+    level: float,
+    device: Union[str, torch.device] = "cuda",
+    max_verts: int = 3_000_000,
+    max_tris: int = 3_000_000,
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Perform the marching cubes algorithm on a 3D grid with warp support.
+    Args:
+        grid_logits (torch.Tensor): A 3D tensor containing the grid logits.
+        level (float): The threshold level for the isosurface.
+        device (Union[str, torch.device], optional): The device to perform the computation on. Defaults to "cuda".
+        max_verts (int, optional): The maximum number of vertices. Defaults to 3,000,000.
+        max_tris (int, optional): The maximum number of triangles. Defaults to 3,000,000.
+    Returns:
+        Tuple[np.ndarray, np.ndarray]: A tuple containing the vertices and faces of the isosurface.
+    """
+    if isinstance(device, torch.device):
+        device = str(device)
+
+    assert grid_logits.ndim == 3
+    if "cuda" in device:
+        assert wp.is_cuda_available()
+    else:
+        raise ValueError(
+            f"Device {device} is not supported for marching_cubes_with_warp"
+        )
+
+    dim = grid_logits.shape[0]
+    field = wp.from_torch(grid_logits)
+
+    iso = wp.MarchingCubes(
+        nx=dim,
+        ny=dim,
+        nz=dim,
+        max_verts=int(max_verts),
+        max_tris=int(max_tris),
+        device=device,
+    )
+    iso.surface(field=field, threshold=level)
+    vertices = iso.verts.numpy()
+    faces = iso.indices.numpy().reshape(-1, 3)
+    return vertices, faces

cube/cube3d/model/autoencoder/one_d_autoencoder.py

@@ -0,0 +1,694 @@
+import logging
+import sys
+from dataclasses import dataclass, field
+from functools import partial
+from typing import List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from skimage import measure
+from torch.nn import functional as F
+from tqdm import tqdm
+
+from cube3d.model.autoencoder.embedder import PhaseModulatedFourierEmbedder
+from cube3d.model.autoencoder.grid import (
+    generate_dense_grid_points,
+    marching_cubes_with_warp,
+)
+from cube3d.model.autoencoder.spherical_vq import SphericalVectorQuantizer
+from cube3d.model.transformers.attention import (
+    EncoderCrossAttentionLayer,
+    EncoderLayer,
+    init_linear,
+    init_tfixup,
+)
+from cube3d.model.transformers.norm import LayerNorm
+
+
+def init_sort(x):
+    """
+    Sorts the input tensor `x` based on its pairwise distances to the first element.
+    This function computes the pairwise distances between all elements in `x` and the
+    first element of `x`. It then sorts the elements of `x` in ascending order of
+    their distances to the first element.
+    Args:
+        x (torch.Tensor): A 2D tensor where each row represents a data point.
+    Returns:
+        torch.Tensor: A tensor containing the rows of `x` sorted by their distances
+        to the first row of `x`.
+    """
+
+    distances = torch.cdist(x, x[:1])
+    _, indices = torch.sort(distances.squeeze(), dim=0)
+    x = x[indices]
+    return x
+
+
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, embed_dim: int, bias: bool = True):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, embed_dim, bias=bias)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        self.apply(partial(init_linear, embed_dim=embed_dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class OneDEncoder(nn.Module):
+    def __init__(
+        self,
+        embedder,
+        num_latents: int,
+        point_feats: int,
+        embed_point_feats: bool,
+        width: int,
+        num_heads: int,
+        num_layers: int,
+        with_cls_token: bool = False,
+        cross_attention_levels: Optional[List[int]] = None,
+        eps: float = 1e-6,
+    ) -> None:
+        """
+        Initializes the OneDEncoder model.
+        Args:
+            embedder: An embedding module that provides the input embedding functionality.
+            num_latents (int): The number of latent variables.
+            point_feats (int): The number of point features.
+            embed_point_feats (bool): Whether to embed point features or not.
+            width (int): The width of the embedding dimension.
+            num_heads (int): The number of attention heads.
+            num_layers (int): The number of encoder layers.
+            with_cls_token (bool, optional): Whether to include a classification token like in Vision Transformers (ViT). Defaults to False.
+            cross_attention_levels (Optional[List[int]], optional): The indices of layers where cross-attention is applied. Defaults to None.
+            eps (float, optional): A small value added for numerical stability in normalization layers. Defaults to 1e-6.
+        Returns:
+            None
+        """
+        super().__init__()
+
+        self.embedder = embedder
+
+        # add cls token like ViT
+        self.with_cls_token = with_cls_token
+        if self.with_cls_token:
+            query = torch.empty((1 + num_latents, width))
+        else:
+            query = torch.empty((num_latents, width))
+
+        # initialize then sort query to potentially get better ordering
+        query.uniform_(-1.0, 1.0)
+        query = init_sort(query)
+
+        # set parameter
+        self.query = nn.Parameter(query)
+
+        self.embed_point_feats = embed_point_feats
+        in_dim = (
+            self.embedder.out_dim * 2
+            if self.embed_point_feats
+            else self.embedder.out_dim + point_feats
+        )
+        self.feat_in = MLPEmbedder(in_dim, embed_dim=width)
+
+        if cross_attention_levels is None:
+            cross_attention_levels = [0]
+
+        self.blocks = nn.ModuleList()
+        for i in range(num_layers):
+            if i in cross_attention_levels:
+                self.blocks.append(
+                    EncoderCrossAttentionLayer(
+                        embed_dim=width,
+                        num_heads=num_heads,
+                        eps=eps,
+                    )
+                )
+            else:
+                self.blocks.append(
+                    EncoderLayer(embed_dim=width, num_heads=num_heads, eps=eps)
+                )
+        self.ln_f = LayerNorm(width, eps=eps)
+
+        init_tfixup(self, num_layers)
+
+    def _forward(self, h, data, attn_mask=None):
+        """
+        Forward pass for the autoencoder model.
+
+        Args:
+            h (torch.Tensor): The input tensor to be processed, typically representing
+                the hidden state or intermediate representation.
+            data (torch.Tensor): The input data tensor to be transformed by the feature
+                extraction layer and used in cross-attention layers.
+            attn_mask (torch.Tensor, optional): An optional attention mask tensor to be
+                used in attention layers for masking specific positions. Defaults to None.
+        Returns:
+            torch.Tensor: The output tensor after processing through the layers and
+            applying final normalization.
+        """
+
+        data = self.feat_in(data)
+
+        for block in self.blocks:
+            if isinstance(block, EncoderCrossAttentionLayer):
+                h = block(h, data)
+            else:
+                h = block(h, attn_mask=attn_mask)
+
+        h = self.ln_f(h)
+        return h
+
+    def forward(
+        self, pts: torch.Tensor, feats: torch.Tensor
+    ) -> Tuple[torch.Tensor, list[torch.Tensor]]:
+        """
+        Forward pass of the 1D autoencoder model.
+        Args:
+            pts (torch.Tensor): Input tensor representing points with shape (batch_size, num_points, point_dim).
+            feats (torch.Tensor): Input tensor representing features with shape (batch_size, num_points, feature_dim).
+                                  Can be None if no features are provided.
+        Returns:
+            Tuple[torch.Tensor, list[torch.Tensor]]:
+                - The output tensor after processing the input data.
+                - A list of intermediate tensors (if applicable) generated during the forward pass.
+        """
+
+        b = pts.shape[0]
+        data = self.embedder(pts)
+
+        if feats is not None:
+            if self.embed_point_feats:
+                feats = self.embedder(feats)
+            data = torch.cat([data, feats], dim=-1)
+
+        # prepare query and data
+        h = self.query.unsqueeze(0).expand(b, -1, -1)
+        return self._forward(h, data, attn_mask=None)
+
+
+class OneDBottleNeck(nn.Module):
+    def __init__(
+        self,
+        block,
+    ) -> None:
+        """
+        Initializes the OneDBottleNeck class.
+        Args:
+            block: The building block or module used within the autoencoder.
+        """
+        super().__init__()
+
+        self.block = block
+
+    def forward(self, h: torch.Tensor) -> Tuple[torch.Tensor, dict]:
+        """
+        Forward pass of the OneDBottleNeck function.
+        Args:
+            h (torch.Tensor): Input tensor to the model.
+        Returns:
+            Tuple[torch.Tensor, dict]: A tuple containing:
+                - The transformed tensor `z` after passing through the block (if applicable).
+                - A dictionary `ret_dict` containing additional information:
+                    - "indices": Indices from the block output (if present).
+                    - "z_q": Quantized tensor from the block output (if present).
+
+        """
+
+        z = h
+        ret_dict = {}
+        if self.block is not None:
+            z, d = self.block(z)
+
+            key_mappings = {
+                "q": "indices",
+                "z_q": "z_q",
+            }
+            for in_key, out_key in key_mappings.items():
+                if in_key in d:
+                    ret_dict[out_key] = d[in_key]
+
+        return z, ret_dict
+
+
+class OneDDecoder(nn.Module):
+    def __init__(
+        self,
+        num_latents: int,
+        width: int,
+        num_heads: int,
+        num_layers: int,
+        eps: float = 1e-6,
+    ) -> None:
+        """
+        Initializes the OneDDecoder class.
+        Args:
+            num_latents (int): The number of latent variables.
+            width (int): The width of the embedding dimension.
+            num_heads (int): The number of attention heads in each encoder layer.
+            num_layers (int): The number of encoder layers.
+            eps (float, optional): A small value added for numerical stability. Defaults to 1e-6.
+        """
+        super().__init__()
+
+        self.register_buffer("query", torch.empty([0, width]), persistent=False)
+        self.positional_encodings = nn.Parameter(
+            init_sort(F.normalize(torch.empty(num_latents, width).normal_()))
+        )
+        self.blocks = nn.ModuleList(
+            [
+                EncoderLayer(embed_dim=width, num_heads=num_heads, eps=eps)
+                for _ in range(num_layers)
+            ]
+        )
+
+        init_tfixup(self, num_layers)
+
+    def _forward(self, h):
+        """
+        Applies a sequence of operations to the input tensor `h` using the blocks
+        defined in the model.
+        Args:
+            h (torch.Tensor): The input tensor to be processed by the blocks.
+        Returns:
+            torch.Tensor: The output tensor after applying all blocks sequentially.
+        """
+
+        for block in self.blocks:
+            h = block(h)
+        return h
+
+    def forward(self, z):
+        """
+        This method processes the input tensor `z` by padding it to a fixed length,
+        adding positional encodings, and then passing it through the `_forward` method.
+
+        Args:
+            z (torch.Tensor): Input tensor.
+        Returns:
+            torch.Tensor: Output tensor after processing through the autoencoder.
+        Notes:
+            - If the `query` attribute has a non-zero shape, the input tensor `z` is padded
+              to match the required length using slices of `query`.
+            - Positional encodings are added to the padded input tensor before passing it
+              to the `_forward` method.
+        """
+
+        # pad input to fixed length
+        if self.query.shape[0] > 0:
+            pad_len = self.query.shape[0] + 1 - z.shape[1]
+            paddings = self.query[:pad_len, ...].unsqueeze(0).expand(z.shape[0], -1, -1)
+            z = torch.cat([paddings, z], dim=1)
+        h = z + self.positional_encodings[: z.shape[1], :].unsqueeze(0).expand(
+            z.shape[0], -1, -1
+        )
+
+        return self._forward(h)
+
+
+class OneDOccupancyDecoder(nn.Module):
+    def __init__(
+        self, embedder, out_features: int, width: int, num_heads: int, eps=1e-6
+    ) -> None:
+        """
+        Initializes the OneDOccupancyDecoder module.
+        Args:
+            embedder: An embedding module that provides input embeddings.
+            out_features (int): The number of output features for the final linear layer.
+            width (int): The width of the intermediate layers.
+            num_heads (int): The number of attention heads for the cross-attention layer.
+            eps (float, optional): A small value added for numerical stability in layer normalization. Defaults to 1e-6.
+        """
+        super().__init__()
+
+        self.embedder = embedder
+        self.query_in = MLPEmbedder(self.embedder.out_dim, width)
+
+        self.attn_out = EncoderCrossAttentionLayer(embed_dim=width, num_heads=num_heads)
+        self.ln_f = LayerNorm(width, eps=eps)
+        self.c_head = nn.Linear(width, out_features)
+
+    def query(self, queries: torch.Tensor):
+        """
+        Processes the input tensor through the embedder and query_in layers.
+        Args:
+            queries (torch.Tensor): A tensor containing the input data to be processed.
+        Returns:
+            torch.Tensor: The output tensor after being processed by the embedder and query_in layers.
+        """
+
+        return self.query_in(self.embedder(queries))
+
+    def forward(self, queries: torch.Tensor, latents: torch.Tensor):
+        """
+        Defines the forward pass of the model.
+        Args:
+            queries (torch.Tensor): Input tensor representing the queries.
+            latents (torch.Tensor): Input tensor representing the latent representations.
+        Returns:
+            torch.Tensor: Output tensor after applying the query transformation,
+                          attention mechanism, and final processing layers.
+        """
+        queries = self.query(queries)
+        x = self.attn_out(queries, latents)
+        x = self.c_head(self.ln_f(x))
+        return x
+
+
+class OneDAutoEncoder(nn.Module):
+    @dataclass
+    class Config:
+        checkpoint_path: str = ""
+
+        # network params
+        num_encoder_latents: int = 256
+        num_decoder_latents: int = 256
+        embed_dim: int = 12
+        width: int = 768
+        num_heads: int = 12
+        out_dim: int = 1
+        eps: float = 1e-6
+
+        # grid features embedding
+        num_freqs: int = 128
+        point_feats: int = 0
+        embed_point_feats: bool = False
+
+        num_encoder_layers: int = 1
+        encoder_cross_attention_levels: list[int] = field(default_factory=list)
+        num_decoder_layers: int = 23
+
+        encoder_with_cls_token: bool = True
+        num_codes: int = 16384
+
+    def __init__(self, cfg: Config) -> None:
+        """
+        Initializes the OneDAutoencoder model.
+        Args:
+            cfg (Config): Configuration object containing the parameters for the model.
+        Attributes:
+            cfg (Config): Stores the configuration object.
+            embedder (PhaseModulatedFourierEmbedder): Embeds input data using phase-modulated Fourier features.
+            encoder (OneDEncoder): Encodes the input data into latent representations.
+            bottleneck (OneDBottleNeck): Bottleneck layer containing a spherical vector quantizer for dimensionality reduction.
+            decoder (OneDDecoder): Decodes latent representations back into the original data space.
+            occupancy_decoder (OneDOccupancyDecoder): Decodes occupancy information from latent representations.
+        """
+
+        super().__init__()
+
+        self.cfg = cfg
+
+        self.embedder = PhaseModulatedFourierEmbedder(
+            num_freqs=self.cfg.num_freqs, input_dim=3
+        )
+
+        self.encoder = OneDEncoder(
+            embedder=self.embedder,
+            num_latents=self.cfg.num_encoder_latents,
+            with_cls_token=self.cfg.encoder_with_cls_token,
+            point_feats=self.cfg.point_feats,
+            embed_point_feats=self.cfg.embed_point_feats,
+            width=self.cfg.width,
+            num_heads=self.cfg.num_heads,
+            num_layers=self.cfg.num_encoder_layers,
+            cross_attention_levels=self.cfg.encoder_cross_attention_levels,
+            eps=self.cfg.eps,
+        )
+
+        block = SphericalVectorQuantizer(
+            self.cfg.embed_dim,
+            self.cfg.num_codes,
+            self.cfg.width,
+            codebook_regularization="kl",
+        )
+        self.bottleneck = OneDBottleNeck(block=block)
+
+        self.decoder = OneDDecoder(
+            num_latents=self.cfg.num_encoder_latents,
+            width=self.cfg.width,
+            num_heads=self.cfg.num_heads,
+            num_layers=self.cfg.num_decoder_layers,
+            eps=self.cfg.eps,
+        )
+
+        self.occupancy_decoder = OneDOccupancyDecoder(
+            embedder=self.embedder,
+            out_features=self.cfg.out_dim,
+            width=self.cfg.width,
+            num_heads=self.cfg.num_heads,
+            eps=self.cfg.eps,
+        )
+
+    @torch.no_grad()
+    def decode_indices(self, shape_ids: torch.Tensor):
+        """
+        Decodes the given shape indices into latent representations.
+        Args:
+            shape_ids (torch.Tensor): A tensor containing the shape indices to be decoded.
+        Returns:
+            torch.Tensor: The decoded latent representations corresponding to the input shape indices.
+        """
+
+        z_q = self.bottleneck.block.lookup_codebook(shape_ids)
+        latents = self.decode(z_q)
+        return latents
+
+    @torch.no_grad()
+    def query_embeds(self, shape_ids: torch.Tensor):
+        """
+        Retrieves the latent embeddings corresponding to the given shape IDs.
+        Args:
+            shape_ids (torch.Tensor): A tensor containing the IDs of the shapes
+                for which the latent embeddings are to be queried.
+        Returns:
+            torch.Tensor: A tensor containing the latent embeddings retrieved
+                from the codebook for the provided shape IDs.
+        """
+
+        z_q = self.bottleneck.block.lookup_codebook_latents(shape_ids)
+        return z_q
+
+    @torch.no_grad()
+    def query_indices(self, shape_embs: torch.Tensor):
+        """
+        Queries the indices of the quantized embeddings from the bottleneck layer.
+        Args:
+            shape_embs (torch.Tensor): The input tensor containing shape embeddings
+                to be quantized.
+        Returns:
+            torch.Tensor: A tensor containing the quantized indices.
+        """
+
+        _, ret_dict = self.bottleneck.block.quantize(shape_embs)
+        return ret_dict["q"]
+
+    def encode(self, x: torch.Tensor, **kwargs):
+        """
+        Encodes the input tensor using the encoder and bottleneck layers.
+        Args:
+            x (torch.Tensor): Input tensor with shape (..., N), where the first 3
+                dimensions represent points (pts) and the remaining dimensions
+                represent features (feats).
+            **kwargs: Additional keyword arguments.
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, None, dict]: A tuple containing:
+                - z_e (torch.Tensor): Encoded tensor before bottleneck processing.
+                - z (torch.Tensor): Encoded tensor after bottleneck processing.
+                - None: Placeholder for compatibility with other methods.
+                - d (dict): Dictionary containing additional information, including:
+                    - "z_cls" (torch.Tensor, optional): Class token if
+                      `self.cfg.encoder_with_cls_token` is True.
+        """
+
+        pts, feats = x[..., :3], x[..., 3:]
+        z_e = self.encoder(pts, feats)
+
+        # split class token
+        if self.cfg.encoder_with_cls_token:
+            z_cls = z_e[:, 0, ...]
+            z_e = z_e[:, 1:, ...]
+
+        # quantize or kl
+        z, d = self.bottleneck(z_e)
+
+        if self.cfg.encoder_with_cls_token:
+            d["z_cls"] = z_cls
+        return z_e, z, None, d
+
+    def decode(self, z: torch.Tensor):
+        """
+        Decodes the latent representation `z` using the decoder network.
+        Args:
+            z (torch.Tensor): The latent representation tensor to be decoded.
+        Returns:
+            torch.Tensor: The decoded output tensor.
+        """
+
+        h = self.decoder(z)
+        return h
+
+    def query(self, queries: torch.Tensor, latents: torch.Tensor):
+        """
+        Computes the logits by decoding the given queries and latent representations.
+        Args:
+            queries (torch.Tensor): A tensor containing the query points to be decoded.
+            latents (torch.Tensor): A tensor containing the latent representations corresponding to the queries.
+        Returns:
+            torch.Tensor: A tensor containing the decoded logits for the given queries and latents.
+        """
+
+        logits = self.occupancy_decoder(queries, latents).squeeze(-1)
+        return logits
+
+    def forward(self, surface, queries, **kwargs):
+        """
+        Perform a forward pass through the autoencoder model.
+        Args:
+            surface (torch.Tensor): The input surface tensor to be encoded.
+            queries (torch.Tensor): The query tensor used for generating logits.
+            **kwargs: Additional keyword arguments.
+        Returns:
+            tuple: A tuple containing:
+                - z (torch.Tensor): The latent representation of the input surface.
+                - latents (torch.Tensor): The decoded output from the latent representation.
+                - None: Placeholder for a potential future return value.
+                - logits (torch.Tensor): The logits generated from the queries and latents.
+                - d (torch.Tensor): Additional output from the encoding process.
+        """
+
+        _, z, _, d = self.encode(surface)
+
+        latents = self.decode(z)
+        logits = self.query(queries, latents)
+
+        return z, latents, None, logits, d
+
+    @torch.no_grad()
+    def extract_geometry(
+        self,
+        latents: torch.FloatTensor,
+        bounds: list[float] = [
+            -1.05,
+            -1.05,
+            -1.05,
+            1.05,
+            1.05,
+            1.05,
+        ],
+        resolution_base: float = 9.0,
+        chunk_size: int = 2_000_000,
+        use_warp: bool = False,
+    ):
+        """
+        Extracts 3D geometry from latent representations using a dense grid sampling
+        and marching cubes algorithm.
+        Args:
+            latents (torch.FloatTensor): A tensor of latent representations with shape
+                (batch_size, latent_dim).
+            bounds (list[float], optional): A list of six floats defining the bounding box
+                for the 3D grid in the format [xmin, ymin, zmin, xmax, ymax, zmax].
+                Defaults to [-1.05, -1.05, -1.05, 1.05, 1.05, 1.05].
+            resolution_base (float, optional): The base resolution for the grid. Higher
+                values result in finer grids. Defaults to 9.0.
+            chunk_size (int, optional): The number of grid points to process in a single
+                chunk. Defaults to 2,000,000.
+            use_warp (bool, optional): Whether to use a GPU-accelerated marching cubes
+                implementation. If False, falls back to a CPU implementation. Defaults to False.
+        Returns:
+            tuple:
+                - mesh_v_f (list[tuple]): A list of tuples containing vertices and faces
+                  for each batch element. Each tuple is of the form
+                  (vertices, faces), where:
+                    - vertices (np.ndarray): Array of vertex coordinates with shape
+                      (num_vertices, 3).
+                    - faces (np.ndarray): Array of face indices with shape
+                      (num_faces, 3).
+                  If geometry extraction fails for a batch element, the tuple will be
+                  (None, None).
+                - has_surface (np.ndarray): A boolean array indicating whether a surface
+                  was successfully extracted for each batch element.
+        Raises:
+            Exception: Logs warnings or errors if geometry extraction fails for any
+            batch element or if the marching cubes algorithm encounters issues.
+        """
+        bbox_min = np.array(bounds[0:3])
+        bbox_max = np.array(bounds[3:6])
+        bbox_size = bbox_max - bbox_min
+
+        xyz_samples, grid_size, length = generate_dense_grid_points(
+            bbox_min=bbox_min,
+            bbox_max=bbox_max,
+            resolution_base=resolution_base,
+            indexing="ij",
+        )
+        xyz_samples = torch.FloatTensor(xyz_samples)
+        batch_size = latents.shape[0]
+
+        batch_logits = []
+
+        progress_bar = tqdm(
+            range(0, xyz_samples.shape[0], chunk_size),
+            desc=f"extracting geometry",
+            unit="chunk",
+        )
+        for start in progress_bar:
+            queries = xyz_samples[start : start + chunk_size, :]
+
+            num_queries = queries.shape[0]
+            if start > 0 and num_queries < chunk_size:
+                queries = F.pad(queries, [0, 0, 0, chunk_size - num_queries])
+            batch_queries = queries.unsqueeze(0).expand(batch_size, -1, -1).to(latents)
+
+            logits = self.query(batch_queries, latents)[:, :num_queries]
+            batch_logits.append(logits)
+
+        grid_logits = (
+            torch.cat(batch_logits, dim=1)
+            .detach()
+            .view((batch_size, grid_size[0], grid_size[1], grid_size[2]))
+            .float()
+        )
+
+        mesh_v_f = []
+        has_surface = np.zeros((batch_size,), dtype=np.bool_)
+        for i in range(batch_size):
+            try:
+                warp_success = False
+                if use_warp:
+                    try:
+                        vertices, faces = marching_cubes_with_warp(
+                            grid_logits[i],
+                            level=0.0,
+                            device=grid_logits.device,
+                        )
+                        warp_success = True
+                    except Exception as e:
+                        logging.warning(
+                            f"Warning: error in marching cubes with warp: {e}"
+                        )
+                        warp_success = False  # Fall back to CPU version
+
+                if not warp_success:
+                    logging.warning(
+                        "Warning: falling back to CPU version of marching cubes using skimage measure"
+                    )
+                    vertices, faces, _, _ = measure.marching_cubes(
+                        grid_logits[i].cpu().numpy(), 0, method="lewiner"
+                    )
+
+                vertices = vertices / grid_size * bbox_size + bbox_min
+                faces = faces[:, [2, 1, 0]]
+                mesh_v_f.append(
+                    (vertices.astype(np.float32), np.ascontiguousarray(faces))
+                )
+                has_surface[i] = True
+            except Exception as e:
+                logging.error(f"Error: error in extract_geometry: {e}")
+                mesh_v_f.append((None, None))
+                has_surface[i] = False
+
+        return mesh_v_f, has_surface

cube/cube3d/model/autoencoder/spherical_vq.py

@@ -0,0 +1,159 @@
+import sys
+from typing import Literal, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cube3d.model.transformers.norm import RMSNorm
+
+
+class SphericalVectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_codes: int,
+        width: Optional[int] = None,
+        codebook_regularization: Literal["batch_norm", "kl"] = "batch_norm",
+    ):
+        """
+        Initializes the SphericalVQ module.
+        Args:
+            embed_dim (int): The dimensionality of the embeddings.
+            num_codes (int): The number of codes in the codebook.
+            width (Optional[int], optional): The width of the input. Defaults to None.
+        Raises:
+            ValueError: If beta is not in the range [0, 1].
+        """
+        super().__init__()
+
+        self.num_codes = num_codes
+
+        self.codebook = nn.Embedding(num_codes, embed_dim)
+        self.codebook.weight.data.uniform_(-1.0 / num_codes, 1.0 / num_codes)
+
+        width = width or embed_dim
+        if width != embed_dim:
+            self.c_in = nn.Linear(width, embed_dim)
+            self.c_x = nn.Linear(width, embed_dim)  # shortcut
+            self.c_out = nn.Linear(embed_dim, width)
+        else:
+            self.c_in = self.c_out = self.c_x = nn.Identity()
+
+        self.norm = RMSNorm(embed_dim, elementwise_affine=False)
+        self.cb_reg = codebook_regularization
+        if self.cb_reg == "batch_norm":
+            self.cb_norm = nn.BatchNorm1d(embed_dim, track_running_stats=False)
+        else:
+            self.cb_weight = nn.Parameter(torch.ones([embed_dim]))
+            self.cb_bias = nn.Parameter(torch.zeros([embed_dim]))
+            self.cb_norm = lambda x: x.mul(self.cb_weight).add_(self.cb_bias)
+
+    def get_codebook(self):
+        """
+        Retrieves the normalized codebook weights.
+        This method applies a series of normalization operations to the
+        codebook weights, ensuring they are properly scaled and normalized
+        before being returned.
+        Returns:
+            torch.Tensor: The normalized weights of the codebook.
+        """
+
+        return self.norm(self.cb_norm(self.codebook.weight))
+
+    @torch.no_grad()
+
+    def lookup_codebook(self, q: torch.Tensor):
+        """
+        Perform a lookup in the codebook and process the result.
+        This method takes an input tensor of indices, retrieves the corresponding
+        embeddings from the codebook, and applies a transformation to the retrieved
+        embeddings.
+        Args:
+            q (torch.Tensor): A tensor containing indices to look up in the codebook.
+        Returns:
+            torch.Tensor: The transformed embeddings retrieved from the codebook.
+        """
+
+        # normalize codebook
+        z_q = F.embedding(q, self.get_codebook())
+        z_q = self.c_out(z_q)
+        return z_q
+
+    @torch.no_grad()
+    def lookup_codebook_latents(self, q: torch.Tensor):
+        """
+        Retrieves the latent representations from the codebook corresponding to the given indices.
+        Args:
+            q (torch.Tensor): A tensor containing the indices of the codebook entries to retrieve.
+                              The indices should be integers and correspond to the rows in the codebook.
+        Returns:
+            torch.Tensor: A tensor containing the latent representations retrieved from the codebook.
+                          The shape of the returned tensor depends on the shape of the input indices
+                          and the dimensionality of the codebook entries.
+        """
+
+        # normalize codebook
+        z_q = F.embedding(q, self.get_codebook())
+        return z_q
+
+    def quantize(self, z: torch.Tensor):
+        """
+        Quantizes the latent codes z with the codebook
+
+        Args:
+                z (Tensor): B x ... x F
+        """
+
+        # normalize codebook
+        codebook = self.get_codebook()
+        # the process of finding quantized codes is non differentiable
+        with torch.no_grad():
+            # flatten z
+            z_flat = z.view(-1, z.shape[-1])
+
+            # calculate distance and find the closest code
+            d = torch.cdist(z_flat, codebook)
+            q = torch.argmin(d, dim=1)  # num_ele
+
+        z_q = codebook[q, :].reshape(*z.shape[:-1], -1)
+        q = q.view(*z.shape[:-1])
+
+        return z_q, {"z": z.detach(), "q": q}
+
+    def straight_through_approximation(self, z, z_q):
+        """passed gradient from z_q to z"""
+        z_q = z + (z_q - z).detach()
+        return z_q
+
+    def forward(self, z: torch.Tensor):
+        """
+        Forward pass of the spherical vector quantization autoencoder.
+        Args:
+            z (torch.Tensor): Input tensor of shape (batch_size, ..., feature_dim).
+        Returns:
+            Tuple[torch.Tensor, Dict[str, Any]]:
+                - z_q (torch.Tensor): The quantized output tensor after applying the
+                  straight-through approximation and output projection.
+                - ret_dict (Dict[str, Any]): A dictionary containing additional
+                  information:
+                    - "z_q" (torch.Tensor): Detached quantized tensor.
+                    - "q" (torch.Tensor): Indices of the quantized vectors.
+                    - "perplexity" (torch.Tensor): The perplexity of the quantization,
+                      calculated as the exponential of the negative sum of the
+                      probabilities' log values.
+        """
+
+        with torch.autocast(device_type=z.device.type, enabled=False):
+            # work in full precision
+            z = z.float()
+
+            # project and normalize
+            z_e = self.norm(self.c_in(z))
+            z_q, ret_dict = self.quantize(z_e)
+
+            ret_dict["z_q"] = z_q.detach()
+            z_q = self.straight_through_approximation(z_e, z_q)
+            z_q = self.c_out(z_q)
+
+        return z_q, ret_dict

cube/cube3d/model/gpt/dual_stream_roformer.py

@@ -0,0 +1,279 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import nn
+
+from cube3d.model.transformers.cache import Cache
+from cube3d.model.transformers.dual_stream_attention import (
+    DualStreamDecoderLayerWithRotaryEmbedding,
+)
+from cube3d.model.transformers.norm import LayerNorm
+from cube3d.model.transformers.roformer import DecoderLayerWithRotaryEmbedding
+from cube3d.model.transformers.rope import precompute_freqs_cis
+
+
+class DualStreamRoformer(nn.Module):
+    @dataclass
+    class Config:
+        checkpoint_path: str = ""
+        n_layer: int = 12
+        n_single_layer: int = 0
+        rope_theta: float = 1000
+
+        n_head: int = 16
+        n_embd: int = 2048
+        bias: bool = False  # bias in Linears and LayerNorms
+        eps: float = 1e-6  # Norm eps
+
+        shape_model_vocab_size: int = 4096
+        shape_model_embed_dim: int = 16
+
+        text_model_embed_dim: int = 512
+        use_pooled_text_embed: bool = False
+
+        encoder_with_cls_token: bool = True
+
+    def __init__(self, cfg: Config) -> None:
+        """
+        Initializes the DualStreamRoFormer model.
+        Args:
+            cfg (Config): Configuration object containing model parameters.
+        Attributes:
+            cfg (Config): Stores the configuration object.
+            text_proj (nn.Linear): Linear layer to project text model embeddings to the desired embedding dimension.
+            shape_proj (nn.Linear, optional): Linear layer to project shape model embeddings to the desired embedding
+                dimension
+            vocab_size (int): Vocabulary size for the shape model, including special tokens.
+            shape_bos_id (int): Token ID for the beginning-of-sequence (BOS) token for the shape model.
+            shape_eos_id (int): Token ID for the end-of-sequence (EOS) token for the shape model.
+            padding_id (int): Token ID for the padding token.
+            transformer (nn.ModuleDict): Dictionary containing the following components:
+                - wte (nn.Embedding): Embedding layer for the vocabulary.
+                - dual_blocks (nn.ModuleList): List of dual-stream decoder layers with rotary embeddings.
+                - single_blocks (nn.ModuleList): List of single-stream decoder layers with rotary embeddings.
+                - ln_f (LayerNorm): Layer normalization applied to the final output.
+            lm_head (nn.Linear): Linear layer mapping the final embeddings to the vocabulary size for language modeling.
+        """
+
+        super().__init__()
+
+        self.cfg = cfg
+
+        self.text_proj = nn.Linear(
+            in_features=self.cfg.text_model_embed_dim,
+            out_features=self.cfg.n_embd,
+            bias=self.cfg.bias,
+        )
+
+        self.shape_proj = nn.Linear(self.cfg.shape_model_embed_dim, self.cfg.n_embd)
+
+        self.vocab_size = self.cfg.shape_model_vocab_size
+
+        def add_special_token():
+            token_id = self.vocab_size
+            self.vocab_size += 1
+            return token_id
+
+        self.shape_bos_id = add_special_token()
+        self.shape_eos_id = add_special_token()
+        self.padding_id = add_special_token()
+
+        self.transformer = nn.ModuleDict(
+            dict(
+                wte=nn.Embedding(
+                    self.vocab_size,
+                    self.cfg.n_embd,
+                    padding_idx=self.padding_id,
+                ),
+                dual_blocks=nn.ModuleList(
+                    [
+                        DualStreamDecoderLayerWithRotaryEmbedding.from_config(
+                            self.cfg, cond_pre_only=(i == self.cfg.n_layer - 1)
+                        )
+                        for i in range(self.cfg.n_layer)
+                    ]
+                ),
+                single_blocks=nn.ModuleList(
+                    [
+                        DecoderLayerWithRotaryEmbedding.from_config(self.cfg)
+                        for _ in range(self.cfg.n_single_layer)
+                    ]
+                ),
+                ln_f=LayerNorm(
+                    self.cfg.n_embd, elementwise_affine=False, eps=self.cfg.eps
+                ),
+            )
+        )
+
+        self.lm_head = nn.Linear(self.cfg.n_embd, self.vocab_size, bias=False)
+
+    def encode_text(self, text_embed):
+        """
+        Encodes the given text embeddings by projecting them through a linear transformation.
+        Args:
+            text_embed (torch.Tensor): A tensor representing the text embeddings to be encoded.
+        Returns:
+            torch.Tensor: The projected text embeddings after applying the linear transformation.
+        """
+
+        return self.text_proj(text_embed)
+
+    def encode_token(self, tokens):
+        """
+        Encodes the input tokens using the word token embedding layer of the transformer model.
+        Args:
+            tokens (torch.Tensor): A tensor containing the input tokens to be encoded.
+        Returns:
+            torch.Tensor: A tensor containing the encoded token embeddings.
+        """
+
+        return self.transformer.wte(tokens)
+
+    def init_kv_cache(
+        self,
+        batch_size: int,
+        cond_len: int,
+        max_shape_tokens: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> list[Cache]:
+        """
+        Initializes the key-value cache for the transformer model.
+        This method creates a list of `Cache` objects to store the key and value
+        states for both dual-stream and single-stream transformer blocks. The
+        cache is pre-allocated with zeros and is used to optimize the computation
+        of attention mechanisms during model inference.
+        Args:
+            batch_size (int): The batch size for the input data.
+            cond_len (int): The length of the conditioning sequence.
+            max_shape_tokens (int): The maximum number of tokens in the shape sequence.
+            dtype (torch.dtype): The data type for the tensors (e.g., torch.float32).
+            device (torch.device): The device on which the tensors will be allocated
+                (e.g., torch.device('cuda') or torch.device('cpu')).
+        Returns:
+            list[Cache]: A list of `Cache` objects containing pre-allocated key and
+            value states for each transformer block.
+        """
+        num_heads = self.cfg.n_head
+        max_all_tokens = cond_len + max_shape_tokens
+        per_head_dim = self.cfg.n_embd // num_heads
+
+        kv_cache = [
+            Cache(
+                key_states=torch.zeros(
+                    (batch_size, num_heads, max_all_tokens, per_head_dim),
+                    dtype=dtype,
+                    device=device,
+                ),
+                value_states=torch.zeros(
+                    (batch_size, num_heads, max_all_tokens, per_head_dim),
+                    dtype=dtype,
+                    device=device,
+                ),
+            )
+            for _ in range(len(self.transformer.dual_blocks))
+        ]
+        kv_cache += [
+            Cache(
+                key_states=torch.zeros(
+                    (batch_size, num_heads, max_shape_tokens, per_head_dim),
+                    dtype=dtype,
+                    device=device,
+                ),
+                value_states=torch.zeros(
+                    (batch_size, num_heads, max_shape_tokens, per_head_dim),
+                    dtype=dtype,
+                    device=device,
+                ),
+            )
+            for _ in range(len(self.transformer.single_blocks))
+        ]
+        return kv_cache
+
+    def forward(
+        self,
+        embed: torch.Tensor,
+        cond: torch.Tensor,
+        kv_cache: Optional[list[Cache]] = None,
+        curr_pos_id: Optional[torch.Tensor] = None,
+        decode: bool = False,
+    ):
+        """
+        Forward pass for the dual-stream RoFormer model.
+        Args:
+            embed (torch.Tensor): The input embedding tensor.
+            cond (torch.Tensor): The conditioning tensor.
+            kv_cache (Optional[list[Cache]]): A list of key-value caches for each layer, used for decoding. Default is None.
+            curr_pos_id (Optional[torch.Tensor]): The current position ID tensor of shape (batch_size,). Required if `decode` is True. Default is None.
+            decode (bool): Whether the model is in decoding mode. Default is False.
+        Returns:
+            torch.Tensor: The output logits tensor.
+        """
+        b, l = embed.shape[:2]
+        s = cond.shape[1]
+        device = embed.device
+
+        attn_mask = torch.tril(
+            torch.ones(s + l, s + l, dtype=torch.bool, device=device)
+        )
+
+        position_ids = torch.arange(l, dtype=torch.long, device=device)  # shape (t)
+        position_ids = position_ids.unsqueeze_(0).expand(b, -1)
+
+        s_freqs_cis = precompute_freqs_cis(
+            dim=self.cfg.n_embd // self.cfg.n_head,
+            t=position_ids,
+            theta=self.cfg.rope_theta,
+        )
+
+        position_ids = torch.cat(
+            [
+                torch.zeros([b, s], dtype=torch.long, device=position_ids.device),
+                position_ids,
+            ],
+            dim=1,
+        )
+        d_freqs_cis = precompute_freqs_cis(
+            dim=self.cfg.n_embd // self.cfg.n_head,
+            t=position_ids,
+            theta=self.cfg.rope_theta,
+        )
+
+        if kv_cache is not None and decode:
+            assert curr_pos_id is not None
+            embed = embed[:, curr_pos_id, :]
+
+        h = embed
+        c = cond
+
+        layer_idx = 0
+        for block in self.transformer.dual_blocks:
+            h, c = block(
+                h,
+                c=c,
+                freqs_cis=d_freqs_cis,
+                attn_mask=attn_mask,
+                is_causal=True,
+                kv_cache=kv_cache[layer_idx] if kv_cache is not None else None,
+                curr_pos_id=curr_pos_id + s if curr_pos_id is not None else None,
+                decode=decode,
+            )
+            layer_idx += 1
+        for block in self.transformer.single_blocks:
+            h = block(
+                h,
+                freqs_cis=s_freqs_cis,
+                attn_mask=None,
+                is_causal=True,
+                kv_cache=kv_cache[layer_idx] if kv_cache is not None else None,
+                curr_pos_id=curr_pos_id,
+                decode=decode,
+            )
+            layer_idx += 1
+
+        # Normalization
+        h = self.transformer.ln_f(h)
+        logits = self.lm_head(h)
+
+        return logits

cube/cube3d/model/transformers/attention.py

@@ -0,0 +1,300 @@
+import math
+
+import torch
+import torch.nn as nn
+
+from cube3d.model.transformers.norm import LayerNorm, RMSNorm
+
+
+def init_linear(module, embed_dim: int):
+    """
+    Initializes the weights and biases of a given linear module.
+    Args:
+        module (nn.Module): The module to initialize. Expected to be an instance of nn.Linear.
+        embed_dim (int): The embedding dimension used to calculate the standard deviation
+                         for weight initialization.
+    Returns:
+        None
+    """
+
+    if isinstance(module, nn.Linear):
+        nn.init.normal_(module.weight, std=math.sqrt(1.0 / embed_dim))
+        if module.bias is not None:
+            torch.nn.init.zeros_(module.bias)
+
+
+def init_tfixup(module: nn.Module, num_layers: int):
+    """Special initialization from https://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf
+
+    Args:
+        module (nn.Module): decoder/encoder module
+        num_layers (int): number of layers in the module
+    """
+    with torch.no_grad():
+        for pn, p in module.named_parameters():
+            if (
+                pn.endswith("c_proj.weight")
+                or pn.endswith("up_proj.weight")
+                or pn.endswith("down_proj.weight")
+            ):
+                p *= (4 * num_layers) ** (-0.25)
+            elif pn.endswith("c_v.weight"):
+                p *= (4 * num_layers) ** (-0.25) * math.sqrt(2)
+
+
+class MLP(nn.Module):
+    def __init__(self, embed_dim, hidden_dim, bias=True, approximate="none"):
+        """
+        MLP with GELU activation function."
+        """
+
+        super().__init__()
+        self.up_proj = nn.Linear(embed_dim, hidden_dim, bias=bias)
+        self.down_proj = nn.Linear(hidden_dim, embed_dim, bias=bias)
+        self.act_fn = nn.GELU(approximate=approximate)
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.up_proj(x)))
+
+
+class SelfAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        bias: bool = True,
+        eps: float = 1e-6,
+    ):
+        """
+        Initializes the self attention mechanism.
+        Args:
+            embed_dim (int): The dimensionality of the embedding space.
+            num_heads (int): The number of attention heads.
+            bias (bool, optional): Whether to include bias terms in the linear layers. Defaults to True.
+            eps (float, optional): A small value added for numerical stability. Defaults to 1e-6.
+        Raises:
+            AssertionError: If `embed_dim` is not divisible by `num_heads`.
+        """
+
+        super().__init__()
+        assert embed_dim % num_heads == 0
+        self.num_heads = num_heads
+        self.c_qk = nn.Linear(embed_dim, 2 * embed_dim, bias=bias)
+        self.c_v = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.c_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        head_dim = embed_dim // num_heads
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+
+    def forward(self, x, attn_mask=None, is_causal: bool = False):
+        """
+        Performs the forward pass of the attention mechanism.
+        Args:
+            x (torch.Tensor): Input tensor.
+            attn_mask (Optional[torch.Tensor]): Attention mask to apply. Default is None.
+            is_causal (bool): If True, applies a causal mask to prevent attending to future positions.
+                              Default is False.
+        Returns:
+            torch.Tensor: Output tensor after applying
+                          the attention mechanism and projection.
+        """
+
+        b, l, d = x.shape
+
+        q, k = self.c_qk(x).chunk(2, dim=-1)
+        v = self.c_v(x)
+
+        q = q.view(b, l, self.num_heads, -1).transpose(1, 2)  # (B, nh, T, hs)
+        k = k.view(b, l, self.num_heads, -1).transpose(1, 2)  # (B, nh, T, hs)
+        v = v.view(b, l, self.num_heads, -1).transpose(1, 2)  # (B, nh, T, hs)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        is_causal = is_causal and attn_mask is None
+        y = torch.nn.functional.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            dropout_p=0.0,
+            is_causal=is_causal,
+        )
+
+        y = y.transpose(1, 2).contiguous().view(b, l, d)
+
+        y = self.c_proj(y)
+        return y
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        q_dim=None,
+        kv_dim=None,
+        bias: bool = True,
+    ):
+        """
+        Initializes the cross attention mechanism.
+        Args:
+            embed_dim (int): The dimensionality of the embedding space.
+            num_heads (int): The number of attention heads.
+            q_dim (int, optional): The dimensionality of the query input. Defaults to `embed_dim`.
+            kv_dim (int, optional): The dimensionality of the key and value inputs. Defaults to `embed_dim`.
+            bias (bool, optional): Whether to include a bias term in the linear projections. Defaults to True.
+        Raises:
+            AssertionError: If `embed_dim` is not divisible by `num_heads`.
+        """
+        super().__init__()
+        assert embed_dim % num_heads == 0
+
+        q_dim = q_dim or embed_dim
+        kv_dim = kv_dim or embed_dim
+
+        self.c_q = nn.Linear(q_dim, embed_dim, bias=bias)
+        self.c_k = nn.Linear(kv_dim, embed_dim, bias=bias)
+        self.c_v = nn.Linear(kv_dim, embed_dim, bias=bias)
+        self.c_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.num_heads = num_heads
+
+    def forward(self, x, c, attn_mask=None, is_causal: bool = False):
+        """
+        Forward pass for the attention mechanism.
+        Args:
+            x (torch.Tensor): Input tensor of shape.
+            c (torch.Tensor): Context tensor.
+            attn_mask (torch.Tensor, optional): Attention mask.
+                Defaults to None.
+            is_causal (bool, optional): Whether to apply causal masking. Defaults to False.
+        Returns:
+            torch.Tensor: Output tensor.
+        """
+
+        q, k = self.c_q(x), self.c_k(c)
+        v = self.c_v(c)
+
+        b, l, d = q.shape
+        s = k.shape[1]
+
+        q = q.view(b, l, self.num_heads, -1).transpose(1, 2)  # (B, nh, T, hs)
+        k = k.view(b, s, self.num_heads, -1).transpose(1, 2)  # (B, nh, T, hs)
+        v = v.view(b, s, self.num_heads, -1).transpose(1, 2)  # (B, nh, T, hs)
+
+        y = torch.nn.functional.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            dropout_p=0.0,
+            is_causal=(attn_mask is not None) and is_causal,
+        )
+
+        y = y.transpose(1, 2).contiguous().view(b, l, d)
+
+        y = self.c_proj(y)
+        return y
+
+
+class EncoderLayer(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        bias: bool = True,
+        eps: float = 1e-6,
+    ) -> None:
+        """
+        Initializes the EncoderLayer module.
+        Args:
+            embed_dim (int): The dimensionality of the embedding space.
+            num_heads (int): The number of attention heads.
+            bias (bool, optional): Whether to include bias terms in the layers. Defaults to True.
+            eps (float, optional): A small value added for numerical stability in normalization layers. Defaults to 1e-6.
+        """
+        super().__init__()
+        self.ln_1 = LayerNorm(embed_dim, elementwise_affine=False, eps=eps)
+        self.attn = SelfAttention(embed_dim, num_heads, bias=bias, eps=eps)
+        self.ln_2 = LayerNorm(embed_dim, elementwise_affine=False, eps=eps)
+        self.mlp = MLP(embed_dim=embed_dim, hidden_dim=embed_dim * 4, bias=bias)
+
+    def forward(self, x, attn_mask=None, is_causal: bool = False):
+        """
+        Performs the forward pass of the transformer block.
+        Args:
+            x (torch.Tensor): The input tensor.
+            attn_mask (torch.Tensor, optional): An optional attention mask tensor to apply during the
+                attention computation. Default is None.
+            is_causal (bool, optional): If True, applies a causal mask to prevent attention to future
+                positions. Default is False.
+        Returns:
+            torch.Tensor: The output tensor of the same shape as the input.
+        """
+
+        x = x + self.attn(self.ln_1(x), attn_mask=attn_mask, is_causal=is_causal)
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class EncoderCrossAttentionLayer(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        q_dim=None,
+        kv_dim=None,
+        bias: bool = True,
+        eps: float = 1e-6,
+    ) -> None:
+        """
+        Initializes the EncoderAttentionLayer module with cross-attention,
+        and a feed-forward MLP.
+        Args:
+            embed_dim (int): The dimensionality of the embedding space.
+            num_heads (int): The number of attention heads.
+            q_dim (int, optional): Dimensionality of the query input. Defaults to `embed_dim`.
+            kv_dim (int, optional): Dimensionality of the key and value inputs. Defaults to `embed_dim`.
+            bias (bool, optional): Whether to include bias terms in the layers. Defaults to True.
+            eps (float, optional): A small value added to the denominator for numerical stability
+                in layer normalization. Defaults to 1e-6.
+        """
+        super().__init__()
+
+        q_dim = q_dim or embed_dim
+        kv_dim = kv_dim or embed_dim
+
+        self.attn = CrossAttention(
+            embed_dim,
+            num_heads,
+            q_dim=q_dim,
+            kv_dim=kv_dim,
+            bias=bias,
+        )
+
+        self.ln_1 = LayerNorm(q_dim, elementwise_affine=False, eps=eps)
+        self.ln_2 = LayerNorm(kv_dim, elementwise_affine=False, eps=eps)
+
+        self.ln_f = LayerNorm(embed_dim, elementwise_affine=False, eps=eps)
+        self.mlp = MLP(embed_dim=embed_dim, hidden_dim=embed_dim * 4, bias=bias)
+
+    def forward(self, x, c, attn_mask=None, is_causal: bool = False):
+        """
+        Forward pass for the attention mechanism.
+        Args:
+            x (torch.Tensor): The input tensor to the attention mechanism.
+            c (torch.Tensor): The context tensor used for cross-attention.
+            attn_mask (torch.Tensor, optional): An optional attention mask to control
+                which positions can attend to others. Defaults to None.
+            is_causal (bool, optional): If True, applies a causal mask to prevent
+                attending to future positions. Defaults to False.
+        Returns:
+            torch.Tensor: The output tensor after applying attention and MLP layers.
+        """
+
+        x = x + self.attn(
+            self.ln_1(x), self.ln_2(c), attn_mask=attn_mask, is_causal=is_causal
+        )
+        x = x + self.mlp(self.ln_f(x))
+        return x

cube/cube3d/model/transformers/cache.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class Cache:
+    key_states: torch.Tensor
+    value_states: torch.Tensor

cube/cube3d/model/transformers/dual_stream_attention.py

@@ -0,0 +1,348 @@
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from cube3d.model.transformers.cache import Cache
+from cube3d.model.transformers.norm import LayerNorm, RMSNorm
+from cube3d.model.transformers.roformer import SwiGLUMLP
+from cube3d.model.transformers.rope import scaled_dot_product_attention_with_rotary_emb
+
+
+class DismantledPreAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        query: bool = True,
+        bias: bool = True,
+    ) -> None:
+        """
+        Initializes the DismantledPreAttention module.
+        Args:
+            embed_dim (int): The dimensionality of the embedding space.
+            num_heads (int): The number of attention heads.
+            query (bool, optional): Whether to include query-key projection. Defaults to True.
+            bias (bool, optional): Whether to include bias in linear layers. Defaults to True.
+        Raises:
+            AssertionError: If `embed_dim` is not divisible by `num_heads`.
+        """
+        super().__init__()
+        assert embed_dim % num_heads == 0
+        self.query = query
+
+        head_dim = embed_dim // num_heads
+        # key, query, value projections for all heads, but in a batch
+        if query:
+            self.c_qk = nn.Linear(embed_dim, 2 * embed_dim, bias=False)
+            self.q_norm = RMSNorm(head_dim)
+        else:
+            self.c_k = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.k_norm = RMSNorm(head_dim)
+        self.c_v = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        # (B, T, C) -> (B, nh, T, hs)
+        self.to_mha = lambda x: x.view(*x.shape[:2], num_heads, -1).transpose(1, 2)
+
+    def forward(self, x):
+        """
+        Forward pass for the dismantled pre-attention mechanism.
+        Args:
+            x (torch.Tensor): Input tensor of shape (..., input_dim).
+        Returns:
+            tuple: A tuple containing:
+                - q (torch.Tensor or None): Query tensor after normalization and transformation,
+                  or None if `self.query` is False.
+                - k (torch.Tensor): Key tensor after normalization and transformation.
+                - v (torch.Tensor): Value tensor after transformation.
+        """
+
+        if self.query:
+            q, k = self.c_qk(x).chunk(2, dim=-1)
+            q = self.q_norm(self.to_mha(q))
+        else:
+            q = None
+            k = self.c_k(x)
+
+        k = self.k_norm(self.to_mha(k))
+        v = self.to_mha(self.c_v(x))
+
+        return (q, k, v)
+
+
+class DismantledPostAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        bias: bool = True,
+        eps: float = 1e-6,
+    ) -> None:
+        """
+        Initializes the DismantledPostAttention module.
+        Args:
+            embed_dim (int): The dimensionality of the embedding space.
+            bias (bool, optional): Whether to include a bias term in the linear projection. Defaults to True.
+            eps (float, optional): A small value added to the denominator for numerical stability in layer normalization. Defaults to 1e-6.
+        """
+        super().__init__()
+        self.c_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.ln_3 = LayerNorm(embed_dim, elementwise_affine=False, eps=eps)
+        self.mlp = SwiGLUMLP(embed_dim, embed_dim * 4, bias=bias)
+
+    def forward(self, x, a):
+        """
+        Forward pass of the dual stream attention mechanism.
+        Args:
+            x (torch.Tensor): The input tensor to the model.
+            a (torch.Tensor): The attention tensor to be combined with the input.
+        Returns:
+            torch.Tensor: The output tensor after applying the projection,
+                          layer normalization, and MLP transformations.
+        """
+
+        x = x + self.c_proj(a)
+        x = x + self.mlp(self.ln_3(x))
+        return x
+
+
+class DualStreamAttentionWithRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        cond_pre_only: bool = False,
+        bias: bool = True,
+    ):
+        """
+        Initializes the DualStreamAttention module.
+        Args:
+            embed_dim (int): The dimensionality of the embedding space.
+            num_heads (int): The number of attention heads.
+            cond_pre_only (bool, optional): If True, the conditional pre-attention
+                will only process the key and value, not the query. Defaults to False.
+            bias (bool, optional): Whether to include a bias term in the attention layers.
+                Defaults to True.
+        """
+        super().__init__()
+
+        self.cond_pre_only = cond_pre_only
+
+        self.pre_x = DismantledPreAttention(
+            embed_dim=embed_dim, num_heads=num_heads, query=True, bias=bias
+        )
+
+        self.pre_c = DismantledPreAttention(
+            embed_dim=embed_dim, num_heads=num_heads, query=not cond_pre_only, bias=bias
+        )
+
+    def forward(
+        self,
+        x,
+        c: Optional[torch.Tensor],
+        freqs_cis,
+        attn_mask: Optional[torch.Tensor] = None,
+        is_causal: bool = False,
+        kv_cache: Optional[Cache] = None,
+        curr_pos_id: Optional[torch.Tensor] = None,
+        decode: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Forward pass for dual stream Multi-Head Attention.
+
+        Efficient single weight matrix multiplication with results split into query, key, value.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Hidden states [B, L, D]
+        c : torch.Tensor
+            Condition [B, S, D]
+        freqs_cis: torch.Tensor
+            Precomputed RoPE matrix from precompute_freqs_cis [B, S+L, Hd]
+        attn_mask : torch.Tensor, optional
+            Attention mask [B, S+L, S+L], by default None
+        kv_cache: None | Tensor
+            key-value cache, but only if not None; if None - it means that it's disabled
+            contains cache for keys and value from all previous steps
+        kv_cache_cond: None | Tensor
+            key-value cache, but only if not None; if None - it means that it's disabled
+            contains cache for keys and value from all previous steps for the text conditioning.
+
+        Returns
+        -------
+        torch.Tensor
+            Hidden state output [B, L, D]
+        """
+        if kv_cache is None or not decode:
+            # Either training or prefill
+            qkv_c = self.pre_c(c)
+            qkv_x = self.pre_x(x)
+            # prepend condition stream
+            # (B, nh, Tc, hs) + (B, nh, Tx, hs) -> (B, nh, Tc+Tx, hs)
+            if self.cond_pre_only:
+                q = qkv_x[0]
+            else:
+                q = torch.cat([qkv_c[0], qkv_x[0]], dim=2)
+            k = torch.cat([qkv_c[1], qkv_x[1]], dim=2)
+            v = torch.cat([qkv_c[2], qkv_x[2]], dim=2)
+
+        else:
+            # if using kv cache, query would only be the last token in the sequence, hence is_causal is False
+            assert x.shape[1] == 1
+            is_causal = False
+            q, k, v = self.pre_x(x)
+
+        if kv_cache is not None:
+            if not decode:
+                kv_cache.key_states[:, :, : k.shape[2], :].copy_(k)
+                kv_cache.value_states[:, :, : k.shape[2], :].copy_(v)
+            else:
+                assert curr_pos_id is not None
+                kv_cache.key_states.index_copy_(2, curr_pos_id, k)
+                kv_cache.value_states.index_copy_(2, curr_pos_id, v)
+            k = kv_cache.key_states
+            v = kv_cache.value_states
+
+        if attn_mask is not None:
+            # trim attention mask to length
+            if decode:
+                assert curr_pos_id is not None
+                attn_mask = attn_mask[..., curr_pos_id, :]
+            else:
+                attn_mask = attn_mask[..., -q.shape[2] :, :]
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        # efficient attention using Flash Attention CUDA kernels
+        y = scaled_dot_product_attention_with_rotary_emb(
+            q,
+            k,
+            v,
+            freqs_cis=freqs_cis,
+            attn_mask=attn_mask,
+            curr_pos_id=curr_pos_id if decode else None,
+            is_causal=is_causal,
+        )
+
+        # re-assemble all head outputs side by side
+        y = y.transpose(1, 2).contiguous().view(x.shape[0], -1, x.shape[2])
+
+        if y.shape[1] == x.shape[1]:
+            y_c = None
+            y_x = y
+        else:
+            assert c is not None, "Conditioning is required for dual stream attention"
+            y_c, y_x = torch.split(y, [c.shape[1], x.shape[1]], dim=1)
+        return y_x, y_c
+
+
+class DualStreamDecoderLayerWithRotaryEmbedding(nn.Module):
+    """Nicely wrapped decoder layer block for dual stream GPT model"""
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads: int,
+        cond_pre_only: bool = False,
+        bias: bool = True,
+        eps: float = 1.0e-6,
+    ) -> None:
+        """
+        Initializes the DualStreamDecoderLayerWithRotaryEmbedding module with optional conditional pre-only mode.
+        Args:
+            embed_dim (int): The dimensionality of the embedding space.
+            num_heads (int): The number of attention heads.
+            cond_pre_only (bool, optional): If True, applies conditional processing only before attention. Defaults to False.
+            bias (bool, optional): If True, includes bias terms in the attention and post-attention layers. Defaults to True.
+            eps (float, optional): A small value added for numerical stability in layer normalization. Defaults to 1.0e-6.
+        """
+        super().__init__()
+
+        self.ln_1 = LayerNorm(embed_dim, elementwise_affine=False, eps=eps)
+        self.ln_2 = LayerNorm(embed_dim, elementwise_affine=False, eps=eps)
+
+        self.attn = DualStreamAttentionWithRotaryEmbedding(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            cond_pre_only=cond_pre_only,
+            bias=bias,
+        )
+
+        self.post_1 = DismantledPostAttention(embed_dim, bias=bias, eps=eps)
+        if not cond_pre_only:
+            self.post_2 = DismantledPostAttention(embed_dim, bias=bias, eps=eps)
+
+    @classmethod
+    def from_config(cls, cfg, cond_pre_only: bool = False):
+        """
+        Create an instance of the class using the provided configuration.
+        Args:
+            cfg: A configuration object containing the necessary parameters:
+                - n_embd (int): The size of the embedding dimension.
+                - n_head (int): The number of attention heads.
+                - bias (bool): Whether to include a bias term.
+                - eps (float): A small value added for numerical stability.
+            cond_pre_only (bool, optional): If True, applies conditioning only in the pre-processing step.
+                Defaults to False.
+        Returns:
+            An instance of the class initialized with the specified configuration.
+        """
+
+        return cls(
+            cfg.n_embd,
+            num_heads=cfg.n_head,
+            cond_pre_only=cond_pre_only,
+            bias=cfg.bias,
+            eps=cfg.eps,
+        )
+
+    def forward(
+        self,
+        x,
+        c,
+        freqs_cis: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        is_causal: bool = True,
+        kv_cache: Optional[Cache] = None,
+        curr_pos_id: Optional[torch.Tensor] = None,
+        decode: bool = False,
+    ):
+        """
+        Forward pass for DualStreamDecoderLayerWithRotaryEmbedding.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Hidden states [B, L, D]
+        c : torch.Tensor
+            Condition [B, S, D]
+        freqs_cis: torch.Tensor
+            Postional embedding from RoPE [B, S+L, hd]
+        attn_mask : torch.Tensor, optional
+            Attention mask [B, S+L, S+L], by default None
+        kv_vache : torch.Tensor, optional
+            kv_cache by default None
+
+        Returns
+        -------
+        torch.Tensor
+            Hidden state output [B, L, D]
+        torch.Tensor
+            kv_cache output [1, L, D]
+        """
+        a_x, a_c = self.attn(
+            self.ln_1(x),
+            # NOTE condition could be none if using kv cache
+            self.ln_2(c) if c is not None else None,
+            freqs_cis=freqs_cis,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+            kv_cache=kv_cache,
+            curr_pos_id=curr_pos_id,
+            decode=decode,
+        )
+        x = self.post_1(x, a_x)
+        if a_c is not None:
+            c = self.post_2(c, a_c)
+        else:
+            c = None
+        return x, c

cube/cube3d/model/transformers/norm.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+
+
+def fused_rms_norm(x: torch.Tensor, weight: nn.Parameter, eps: float):
+    """
+    Applies a fused Root Mean Square (RMS) normalization to the input tensor.
+    Args:
+        x (torch.Tensor): The input tensor to be normalized. Expected to have
+            at least one dimension.
+        weight (nn.Parameter): A learnable parameter used to scale the normalized
+            tensor. Its shape must be broadcastable to the shape of `x`.
+        eps (float): A small constant added to the denominator for numerical
+            stability during normalization.
+    Returns:
+        torch.Tensor: The normalized and scaled tensor with the same shape as `x`.
+    """
+
+    x = x.float()
+    return (x * torch.rsqrt((x * x).mean(-1, keepdim=True).add_(eps))) * weight
+
+
+class LayerNorm(nn.LayerNorm):
+    def forward(self, input: torch.Tensor):
+        """
+        Wrapper to ensure that the input tensor is cast to float before normalization.
+        """
+        y = super().forward(input.float())
+        return y.type_as(input)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine: bool = True):
+        """
+        Initializes the normalization layer.
+        Args:
+            dim (int): The number of features in the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Defaults to 1e-5.
+            elementwise_affine (bool, optional): If True, this layer will have learnable per-element affine parameters. Defaults to True.
+        """
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim), requires_grad=elementwise_affine)
+
+    def forward(self, x):
+        return fused_rms_norm(x, weight=self.weight, eps=self.eps).type_as(x)

cube/cube3d/model/transformers/roformer.py

@@ -0,0 +1,229 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cube3d.model.transformers.cache import Cache
+from cube3d.model.transformers.norm import LayerNorm, RMSNorm
+from cube3d.model.transformers.rope import scaled_dot_product_attention_with_rotary_emb
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, embed_dim, hidden_dim, bias=True, **kwargs):
+        """
+        A PyTorch implementation of the SwiGLU (Swish-Gated Linear Unit) MLP layer.
+        This module consists of three linear projections: `gate_proj`, `up_proj`, and `down_proj`.
+        It applies the SwiGLU activation function, which combines the Swish activation with a gating mechanism,
+        followed by a projection back to the original embedding dimension.
+        Args:
+            embed_dim (int): The dimensionality of the input embeddings.
+            hidden_dim (int): The dimensionality of the hidden layer.
+            bias (bool, optional): Whether to include bias terms in the linear layers. Defaults to True.
+            **kwargs: Additional keyword arguments (currently unused).
+        """
+        super().__init__()
+        self.gate_proj = nn.Linear(embed_dim, hidden_dim, bias=bias)
+        self.up_proj = nn.Linear(embed_dim, hidden_dim, bias=bias)
+        self.down_proj = nn.Linear(hidden_dim, embed_dim, bias=bias)
+
+    # Ignore copy
+    def forward(self, x):
+        """
+        Applies a forward pass.
+        Args:
+            x (torch.Tensor): The input tensor.
+        Returns:
+            torch.Tensor: The output tensor after applying the forward pass.
+        """
+
+        down_proj = self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class SelfAttentionWithRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        bias: bool = True,
+        eps: float = 1e-6,
+    ):
+        """
+        A PyTorch module implementing self-attention with rotary embeddings.
+
+        Args:
+            embed_dim (int): The dimensionality of the input embeddings.
+            num_heads (int): The number of attention heads.
+            bias (bool, optional): Whether to include bias terms in the linear projections. Defaults to True.
+            eps (float, optional): A small value added for numerical stability in normalization. Defaults to 1e-6.
+        """
+        super().__init__()
+        assert embed_dim % num_heads == 0
+        self.num_heads = num_heads
+        # key, query, value projections for all heads, but in a batch
+        self.c_qk = nn.Linear(embed_dim, 2 * embed_dim, bias=False)
+        self.c_v = nn.Linear(embed_dim, embed_dim, bias=bias)
+        # output projection
+        self.c_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        head_dim = embed_dim // num_heads
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+
+    def forward(
+        self,
+        x,
+        freqs_cis: torch.Tensor,
+        attn_mask=None,
+        is_causal: bool = False,
+        kv_cache: Optional[Cache] = None,
+        curr_pos_id: Optional[torch.Tensor] = None,
+        decode: bool = False,
+    ):
+        """
+        Forward pass for the SelfAttentionWithRotaryEmbedding instance.
+        Args:
+            x (torch.Tensor): Input tensor.
+            freqs_cis (torch.Tensor): Precomputed rotary positional embeddings.
+            attn_mask (Optional[torch.Tensor], optional): Attention mask to apply during self-attention. Defaults to None.
+            is_causal (bool, optional): Whether to apply causal masking for autoregressive decoding. Defaults to False.
+            kv_cache (Optional[Cache], optional): Cache object for storing key and value states for decoding. Defaults to None.
+            curr_pos_id (Optional[torch.Tensor], optional): Current position indices for decoding. Required if `decode` is True. Defaults to None.
+            decode (bool, optional): Whether the model is in decoding mode. Defaults to False.
+        Returns:
+            torch.Tensor: Output tensor after applying self-attention and projection.
+        """
+        # batch size, sequence length, embedding dim
+        b, l, d = x.shape
+
+        # compute q, k, v and then split per q, k, v
+        q, k = self.c_qk(x).chunk(2, dim=-1)
+        v = self.c_v(x)
+
+        # split per head
+        q = q.view(b, l, self.num_heads, -1).transpose(1, 2)  # (B, nh, T, hs)
+        k = k.view(b, l, self.num_heads, -1).transpose(1, 2)  # (B, nh, T, hs)
+        v = v.view(b, l, self.num_heads, -1).transpose(1, 2)  # (B, nh, T, hs)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if kv_cache is not None:
+            if not decode:
+                kv_cache.key_states[:, :, : k.shape[2], :].copy_(k)
+                kv_cache.value_states[:, :, : k.shape[2], :].copy_(v)
+            else:
+                assert curr_pos_id is not None
+                kv_cache.key_states.index_copy_(2, curr_pos_id, k)
+                kv_cache.value_states.index_copy_(2, curr_pos_id, v)
+            k = kv_cache.key_states
+            v = kv_cache.value_states
+
+        # self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        # efficient attention using Flash Attention CUDA kernels
+        y = scaled_dot_product_attention_with_rotary_emb(
+            q,
+            k,
+            v,
+            freqs_cis=freqs_cis,
+            attn_mask=attn_mask,
+            curr_pos_id=curr_pos_id if decode else None,
+            is_causal=is_causal,
+        )
+
+        y = (
+            y.transpose(1, 2).contiguous().view(b, l, d)
+        )  # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.c_proj(y)
+        return y
+
+
+class DecoderLayerWithRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        bias: bool = True,
+        eps: float = 1e-6,
+    ) -> None:
+        """
+        Initializes the transformer model with rotary embeddings.
+        Args:
+            embed_dim (int): The dimensionality of the embedding space.
+            num_heads (int): The number of attention heads.
+            bias (bool, optional): Whether to include bias terms in the layers. Defaults to True.
+            eps (float, optional): A small value added for numerical stability in layer normalization. Defaults to 1e-6.
+        """
+        super().__init__()
+
+        self.ln_1 = LayerNorm(embed_dim, elementwise_affine=False, eps=eps)
+        self.attn = SelfAttentionWithRotaryEmbedding(
+            embed_dim, num_heads=num_heads, bias=bias, eps=eps
+        )
+        self.ln_2 = LayerNorm(embed_dim, elementwise_affine=False, eps=eps)
+        self.mlp = SwiGLUMLP(embed_dim, embed_dim * 4, bias=bias)
+
+    @classmethod
+    def from_config(cls, cfg):
+        """
+        Create an instance of the class using the provided configuration.
+        Args:
+            cfg: A configuration object containing the following attributes:
+                - n_embd (int): The size of the embedding dimension.
+                - n_head (int): The number of attention heads.
+                - bias (bool): Whether to include a bias term.
+                - eps (float): A small value added for numerical stability.
+        Returns:
+            An instance of the class initialized with the specified configuration.
+        """
+
+        return cls(
+            cfg.n_embd,
+            num_heads=cfg.n_head,
+            bias=cfg.bias,
+            eps=cfg.eps,
+        )
+
+    def forward(
+        self,
+        x,
+        freqs_cis: torch.Tensor,
+        attn_mask=None,
+        is_causal: bool = True,
+        kv_cache: Optional[Cache] = None,
+        curr_pos_id: Optional[torch.Tensor] = None,
+        decode: bool = False,
+    ):
+        """
+        Forward pass for the transformer model.
+        Args:
+            x (torch.Tensor): Input tensor.
+            freqs_cis (torch.Tensor): Precomputed sinusoidal positional encodings.
+            attn_mask (Optional[torch.Tensor], optional): Attention mask to apply during self-attention.
+                Defaults to None.
+            is_causal (bool, optional): Whether to apply causal masking for autoregressive decoding.
+                Defaults to True.
+            kv_cache (Optional[Cache], optional): Key-value cache for efficient decoding.
+                Defaults to None.
+            curr_pos_id (Optional[torch.Tensor], optional): Current position IDs for decoding.
+                Defaults to None.
+            decode (bool, optional): Whether the model is in decoding mode.
+                Defaults to False.
+        Returns:
+            torch.Tensor: Output tensor.
+        """
+        out = self.attn(
+            self.ln_1(x),
+            freqs_cis=freqs_cis,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+            kv_cache=kv_cache,
+            curr_pos_id=curr_pos_id,
+            decode=decode,
+        )
+        x = x + out
+        x = x + self.mlp(self.ln_2(x))
+        return x

cube/cube3d/model/transformers/rope.py

@@ -0,0 +1,91 @@
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+
+def apply_rotary_emb(
+    x: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    curr_pos_id: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Applies rotary positional embeddings to the input tensor.
+    Args:
+        x (torch.Tensor): The input tensor.
+        freqs_cis (torch.Tensor): A tensor containing the precomputed rotary
+            frequency components.
+        curr_pos_id (Optional[torch.Tensor]): An optional tensor specifying the
+            current position IDs to use for selecting a subset of `freqs_cis`.
+            If None, the function uses the last `seq_len` positions.
+    Returns:
+        torch.Tensor: The input tensor `x` with rotary positional embeddings
+        applied.
+    """
+    x_ = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    if curr_pos_id is None:
+        freqs_cis = freqs_cis[:, -x.shape[2] :].unsqueeze(1)
+    else:
+        freqs_cis = freqs_cis[:, curr_pos_id, :].unsqueeze(1)
+    y = torch.view_as_real(x_ * freqs_cis).flatten(3)
+    return y.type_as(x)
+
+
+@torch.no_grad
+def precompute_freqs_cis(dim: int, t: torch.Tensor, theta: float = 10000.0):
+    """Calculate rotary embedding cos & sin, this is useful when every blocks in the network use same positional embedding.
+
+    Args:
+        dim (int): dimension of the single head of the transformer block
+        t (torch.Tensor): position ids [..., L]
+        theta (int, optional): rope theta. Defaults to 10000.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: tuple of cos and sin of rope
+    """
+    assert dim % 2 == 0, (
+        "RoPE only supports embedding dimensions that are multiples of 2"
+    )
+    freqs = 1.0 / (
+        theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=t.device) / dim)
+    )
+    # [batch_size, seq_len, num_freqs]
+    freqs = torch.outer(t.contiguous().view(-1), freqs).reshape(*t.shape, -1)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+
+    return freqs_cis
+
+
+def scaled_dot_product_attention_with_rotary_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    curr_pos_id: Optional[torch.Tensor] = None,
+    is_causal: bool = False,
+) -> torch.Tensor:
+    """
+    Computes scaled dot product attention on query, key and value tensors
+    with rotary position embeddings on query and key.
+
+    Without caching enabled,
+        q should be (bs, nh, seqlen, hd).
+        k and v should stay unchanged, (bs, nh, seqlen, hd).
+    With caching enabled,
+        q should be (bs, nh, 1, hd).
+        k and v should stay unchanged, (bs, nh, 1, hd).
+        causal_mask must be False.
+    """
+    q = apply_rotary_emb(q, freqs_cis, curr_pos_id=curr_pos_id)  # (bs, nh, l, hd)
+    k = apply_rotary_emb(k, freqs_cis, curr_pos_id=None)  # (bs, nh, s + l, hd)
+
+    x = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=attn_mask,
+        dropout_p=0.0,
+        is_causal=is_causal and attn_mask is None,
+    )
+    return x

cube/cube3d/renderer/blender_script.py

@@ -0,0 +1,723 @@
+"""
+Blender script to render images of 3D models.
+
+This script is adopted from the Trellis rendering script:
+https://github.com/microsoft/TRELLIS/blob/main/dataset_toolkits/render.py
+
+"""
+
+import argparse
+import math
+import os
+import platform
+import random
+import sys
+from pathlib import Path
+from typing import Any, Callable, Dict, Generator, Literal, Optional, Tuple
+
+import bpy
+import numpy as np
+from mathutils import Vector
+
+pathdir = Path(__file__).parent
+sys.path.append(pathdir.as_posix())
+
+print(dir(bpy), bpy.__path__)
+
+IMPORT_FUNCTIONS: Dict[str, Callable] = {
+    ".obj": bpy.ops.wm.obj_import,
+    ".glb": bpy.ops.import_scene.gltf,
+    ".gltf": bpy.ops.import_scene.gltf,
+}
+
+
+def center_and_scale_mesh(scale_value: float = 1.0) -> None:
+    """Centers and scales the scene to fit in a unit cube.
+    For example,
+        scale_value = 1.0 ==> [-0.5, 0.5]
+        scale_value = 2.0 ==> [-1.0, 1.0]
+    """
+    # Get all mesh objects
+    mesh_objects = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+    if not mesh_objects:
+        return
+
+    # Calculate bounds
+    min_coords = Vector((float("inf"),) * 3)
+    max_coords = Vector((float("-inf"),) * 3)
+
+    for obj in mesh_objects:
+        # Get all vertices in world space
+        for vertex in obj.data.vertices:
+            world_coord = obj.matrix_world @ vertex.co
+            min_coords.x = min(min_coords.x, world_coord.x)
+            min_coords.y = min(min_coords.y, world_coord.y)
+            min_coords.z = min(min_coords.z, world_coord.z)
+            max_coords.x = max(max_coords.x, world_coord.x)
+            max_coords.y = max(max_coords.y, world_coord.y)
+            max_coords.z = max(max_coords.z, world_coord.z)
+
+    # Calculate center and dimensions
+    center = (min_coords + max_coords) / 2
+    dimensions = max_coords - min_coords
+    scale = scale_value / max(
+        dimensions.x, dimensions.y, dimensions.z
+    )  # Scale to fit in [-scale_value/2, scale_value/2] cube
+
+    # Create an empty to serve as the parent
+    empty = bpy.data.objects.new("Parent_Empty", None)
+    bpy.context.scene.collection.objects.link(empty)
+
+    # Parent all mesh objects to the empty
+    for obj in mesh_objects:
+        obj.parent = empty
+
+    # Move empty to center everything
+    empty.location = -center
+
+    # Apply scale to empty
+    empty.scale = (scale, scale, scale)
+
+    bpy.context.view_layer.update()
+    bpy.ops.object.select_all(action="DESELECT")
+    empty.select_set(True)
+    bpy.context.view_layer.objects.active = empty
+    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
+    print(f"Empty location: {empty.location}")
+    print(f"Empty scale: {empty.scale}")
+
+    return scale
+
+
+def normalize_scene() -> None:
+    """Normalizes the scene by scaling and translating it to fit in a unit cube centered
+    at the origin.
+
+    Mostly taken from the Point-E / Shap-E rendering script
+    (https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
+    but fix for multiple root objects: (see bug report here:
+    https://github.com/openai/shap-e/pull/60).
+
+    Returns:
+        The new parent object that all objects descend from.
+    """
+    if len(list(get_scene_root_objects())) > 1:
+        # create an empty object to be used as a parent for all root objects
+        parent_empty = bpy.data.objects.new("ParentEmpty", None)
+        bpy.context.scene.collection.objects.link(parent_empty)
+
+        # parent all root objects to the empty object
+        for obj in get_scene_root_objects():
+            if obj != parent_empty:
+                obj.parent = parent_empty
+
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    for obj in get_scene_root_objects():
+        obj.scale = obj.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    for obj in get_scene_root_objects():
+        obj.matrix_world.translation += offset
+    bpy.ops.object.select_all(action="DESELECT")
+    bbox_min, bbox_max = scene_bbox()
+    print(f"After normalize_scene: bbox_min: {bbox_min}, bbox_max: {bbox_max}")
+
+    # unparent the camera
+    bpy.data.objects["Camera"].parent = None
+
+    return parent_empty
+
+
+def reset_cameras() -> None:
+    """Resets the cameras in the scene to a single default camera."""
+    # Delete all existing cameras
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.ops.object.select_by_type(type="CAMERA")
+    bpy.ops.object.delete()
+
+    # Create a new camera with default properties
+    bpy.ops.object.camera_add()
+
+    # Rename the new camera to 'NewDefaultCamera'
+    new_camera = bpy.context.active_object
+    new_camera.name = "Camera"
+
+    # Set the new camera as the active camera for the scene
+    scene.camera = new_camera
+
+
+def get_camera_with_position(x, y, z, fov_degrees=40):
+    camera = bpy.data.objects["Camera"]
+    camera.data.angle = math.radians(fov_degrees)
+    camera.location = np.array([x, y, z])
+    direction = -camera.location
+    rot_quat = direction.to_track_quat("-Z", "Y")
+    camera.rotation_euler = rot_quat.to_euler()
+    return camera
+
+
+def reset_scene() -> None:
+    """Resets the scene to a clean state.
+
+    Returns:
+        None
+    """
+    # delete everything that isn't part of a camera or a light
+    for obj in bpy.data.objects:
+        if obj.type not in {"CAMERA", "LIGHT"}:
+            bpy.data.objects.remove(obj, do_unlink=True)
+
+    # delete all the materials
+    for material in bpy.data.materials:
+        bpy.data.materials.remove(material, do_unlink=True)
+
+    # delete all the textures
+    for texture in bpy.data.textures:
+        bpy.data.textures.remove(texture, do_unlink=True)
+
+    # delete all the images
+    for image in bpy.data.images:
+        bpy.data.images.remove(image, do_unlink=True)
+
+
+def load_object(object_path: str) -> None:
+    """Loads a model with a supported file extension into the scene.
+
+    Args:
+        object_path (str): Path to the model file.
+
+    Raises:
+        ValueError: If the file extension is not supported.
+
+    Returns:
+        None
+    """
+    file_extension = Path(object_path).suffix
+    if file_extension is None or file_extension == "":
+        raise ValueError(f"Unsupported file type: {object_path}")
+
+    # load from existing import functions
+    import_function = IMPORT_FUNCTIONS[file_extension]
+
+    if file_extension in {".glb", ".gltf"}:
+        import_function(filepath=object_path, merge_vertices=True)
+    else:
+        import_function(filepath=object_path)
+
+
+def clear_lights():
+    bpy.ops.object.select_all(action="DESELECT")
+    for obj in bpy.context.scene.objects.values():
+        if isinstance(obj.data, bpy.types.Light):
+            obj.select_set(True)
+    bpy.ops.object.delete()
+
+
+def create_light(
+    location,
+    energy=1.0,
+    angle=0.5 * math.pi / 180,
+    light_type: Literal["POINT", "SUN", "SPOT", "AREA"] = "SUN",
+):
+    # https://blender.stackexchange.com/questions/215624/how-to-create-a-light-with-the-python-api-in-blender-2-92
+    light_data = bpy.data.lights.new(name="Light", type=light_type)
+    light_data.energy = energy
+    if light_type != "AREA" and light_type != "POINT":
+        light_data.angle = angle
+    light_object = bpy.data.objects.new(name="Light", object_data=light_data)
+
+    direction = -location
+    rot_quat = direction.to_track_quat("-Z", "Y")
+    light_object.rotation_euler = rot_quat.to_euler()
+    bpy.context.view_layer.update()
+
+    bpy.context.collection.objects.link(light_object)
+    light_object.location = location
+
+
+def create_uniform_lights(
+    distance=2.0,
+    energy=3.0,
+    light_type: Literal["POINT", "SUN", "SPOT", "AREA"] = "SUN",
+):
+    clear_lights()
+    create_light(Vector([1, 0, 0]) * distance, energy=energy, light_type=light_type)
+    create_light(-Vector([1, 0, 0]) * distance, energy=energy, light_type=light_type)
+    create_light(Vector([0, 1, 0]) * distance, energy=energy, light_type=light_type)
+    create_light(-Vector([0, 1, 0]) * distance, energy=energy, light_type=light_type)
+    create_light(Vector([0, 0, 1]) * distance, energy=energy, light_type=light_type)
+    create_light(-Vector([0, 0, 1]) * distance, energy=energy, light_type=light_type)
+
+
+def create_light_at_camera_position(
+    camera_position: Vector,
+    energy=1.5,
+    use_shadow=False,
+    light_type: Literal["POINT", "SUN", "SPOT", "AREA"] = "SUN",
+):
+    clear_lights()
+    create_light(camera_position, energy=energy, light_type=light_type)
+    # disable shadows
+    if not use_shadow:
+        for light in bpy.data.lights:
+            light.use_shadow = False
+
+
+def set_world_background_color(
+    color: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0),
+) -> None:
+    bpy.context.scene.world.use_nodes = True
+    bpy.context.scene.world.node_tree.nodes["Background"].inputs[
+        0
+    ].default_value = color
+    bpy.context.scene.view_settings.view_transform = "Standard"
+
+
+def scene_bbox(
+    single_obj: Optional[bpy.types.Object] = None, ignore_matrix: bool = False
+) -> Tuple[Vector, Vector]:
+    """Returns the bounding box of the scene.
+
+    Taken from Shap-E rendering script
+    (https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
+
+    Args:
+        single_obj (Optional[bpy.types.Object], optional): If not None, only computes
+            the bounding box for the given object. Defaults to None.
+        ignore_matrix (bool, optional): Whether to ignore the object's matrix. Defaults
+            to False.
+
+    Raises:
+        RuntimeError: If there are no objects in the scene.
+
+    Returns:
+        Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
+    """
+    bbox_min = (math.inf,) * 3
+    bbox_max = (-math.inf,) * 3
+    found = False
+    for obj in get_scene_meshes() if single_obj is None else [single_obj]:
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            if not ignore_matrix:
+                coord = obj.matrix_world @ coord
+            bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+            bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+
+    if not found:
+        raise RuntimeError("no objects in scene to compute bounding box for")
+
+    return Vector(bbox_min), Vector(bbox_max)
+
+
+def get_scene_root_objects() -> Generator[bpy.types.Object, None, None]:
+    """Returns all root objects in the scene.
+
+    Yields:
+        Generator[bpy.types.Object, None, None]: Generator of all root objects in the
+            scene.
+    """
+    for obj in bpy.context.scene.objects.values():
+        if not obj.parent and not isinstance(obj.data, bpy.types.Light):
+            yield obj
+
+
+def get_scene_meshes() -> Generator[bpy.types.Object, None, None]:
+    """Returns all meshes in the scene.
+
+    Yields:
+        Generator[bpy.types.Object, None, None]: Generator of all meshes in the scene.
+    """
+    for obj in bpy.context.scene.objects.values():
+        if isinstance(obj.data, (bpy.types.Mesh)):
+            yield obj
+
+
+def delete_missing_textures() -> Dict[str, Any]:
+    """Deletes all missing textures in the scene.
+
+    Returns:
+        Dict[str, Any]: Dictionary with keys "count", "files", and "file_path_to_color".
+            "count" is the number of missing textures, "files" is a list of the missing
+            texture file paths, and "file_path_to_color" is a dictionary mapping the
+            missing texture file paths to a random color.
+    """
+    missing_file_count = 0
+    out_files = []
+    file_path_to_color = {}
+
+    # Check all materials in the scene
+    for material in bpy.data.materials:
+        if material.use_nodes:
+            for node in material.node_tree.nodes:
+                if node.type == "TEX_IMAGE":
+                    image = node.image
+                    if image is not None:
+                        file_path = bpy.path.abspath(image.filepath)
+                        if file_path == "":
+                            # means it's embedded
+                            continue
+
+                        if not os.path.exists(file_path):
+                            # Find the connected Principled BSDF node
+                            connected_node = node.outputs[0].links[0].to_node
+
+                            if connected_node.type == "BSDF_PRINCIPLED":
+                                if file_path not in file_path_to_color:
+                                    # Set a random color for the unique missing file path
+                                    random_color = [random.random() for _ in range(3)]
+                                    file_path_to_color[file_path] = random_color + [1]
+
+                                connected_node.inputs[
+                                    "Base Color"
+                                ].default_value = file_path_to_color[file_path]
+
+                            # Delete the TEX_IMAGE node
+                            material.node_tree.nodes.remove(node)
+                            missing_file_count += 1
+                            out_files.append(image.filepath)
+    return {
+        "count": missing_file_count,
+        "files": out_files,
+        "file_path_to_color": file_path_to_color,
+    }
+
+
+def setup_environment_lighting(envmap_path):
+    world = bpy.context.scene.world
+    world.use_nodes = True
+    nodes = world.node_tree.nodes
+    links = world.node_tree.links
+
+    # Clear existing nodes
+    for node in nodes:
+        nodes.remove(node)
+
+    # Create Background node
+    bg_node = nodes.new(type="ShaderNodeBackground")
+    bg_node.location = (0, 0)
+
+    # Create Environment Texture node
+    env_tex_node = nodes.new(type="ShaderNodeTexEnvironment")
+    env_tex_node.location = (-300, 0)
+
+    # Set the environment texture path (replace this with your file path)
+    env_tex_node.image = bpy.data.images.load(envmap_path)
+
+    # Create World Output node
+    world_output_node = nodes.new(type="ShaderNodeOutputWorld")
+    world_output_node.location = (300, 0)
+
+    # Link nodes
+    links.new(env_tex_node.outputs["Color"], bg_node.inputs["Color"])
+    links.new(bg_node.outputs["Background"], world_output_node.inputs["Surface"])
+
+
+def create_solid_color_material(name, color):
+    mat = bpy.data.materials.new(name)
+    mat.use_nodes = True
+    node_tree = mat.node_tree
+    color_node = node_tree.nodes.new("ShaderNodeBsdfDiffuse")
+    color_node.inputs["Color"].default_value = color
+    mat_output = node_tree.nodes["Material Output"]
+    node_tree.links.new(color_node.outputs["BSDF"], mat_output.inputs["Surface"])
+    return mat
+
+
+def create_phong_material(name, color):
+    mat = bpy.data.materials.new(name)
+    mat.use_nodes = True
+    node_tree = mat.node_tree
+    spec_node = node_tree.nodes.new("ShaderNodeBsdfPrincipled")
+    print(spec_node.inputs.keys())
+    spec_node.inputs["Base Color"].default_value = color
+    spec_node.inputs["Roughness"].default_value = 0.5
+    spec_node.inputs["Metallic"].default_value = 1.0
+    mat_output = node_tree.nodes["Material Output"]
+    node_tree.links.new(spec_node.outputs["BSDF"], mat_output.inputs["Surface"])
+    return mat
+
+
+def render_object(
+    object_file: str,
+    num_renders: int,
+    output_dir: str,
+    transparent_background: bool = False,
+    environment_map: str = None,
+) -> None:
+    """Saves rendered images for given asset to specified output directory.
+
+    Args:
+        object_file (str): Path to the object file.
+        num_renders (int): Number of renders to save of the object.
+        output_dir (str): Path to the directory where the rendered images and metadata
+            will be saved. The rendered images will be saved in the subdirectory
+            `output_dir/stemname`.
+        transparent_background (bool): Whether to use transparent background,
+            otherwise the background is white.
+    Returns:
+        None
+    """
+    os.makedirs(output_dir, exist_ok=True)
+
+    # load the object
+    reset_scene()
+    load_object(object_file)
+
+    if transparent_background:
+        scene.render.film_transparent = True
+    else:
+        scene.render.film_transparent = False
+
+    set_world_background_color([0.2, 0.2, 0.2, 1.0])
+
+    # normalize the scene
+    _ = normalize_scene()
+
+    # Set up cameras
+    cam = scene.objects["Camera"]
+    fov_degrees = 40.0
+    cam.data.angle = np.radians(fov_degrees)
+
+    # Set up camera constraints
+    cam_constraint = cam.constraints.new(type="TRACK_TO")
+    cam_constraint.track_axis = "TRACK_NEGATIVE_Z"
+    cam_constraint.up_axis = "UP_Y"
+    empty = bpy.data.objects.new("Empty", None)
+    empty.location = (0, 0, 0)
+    scene.collection.objects.link(empty)
+    cam_constraint.target = empty
+    cam.parent = empty
+
+    # delete all objects that are not meshes
+    delete_missing_textures()
+
+    if environment_map:
+        setup_environment_lighting(environment_map)
+    else:
+        create_uniform_lights(energy=1.0, light_type="SUN")
+
+    camera_position = [0, -2, 0]
+
+    # determine how much to orbit camera by.
+    stepsize = 360.0 / num_renders
+
+    def render_views(name):
+        for i in range(num_renders):
+            # set camera
+            _ = get_camera_with_position(
+                camera_position[0],
+                camera_position[1],
+                camera_position[2],
+                fov_degrees=fov_degrees,
+            )
+
+            # Set output paths with absolute paths
+            render_path = os.path.abspath(
+                os.path.join(output_dir, f"{i:03d}_{name}.png")
+            )
+
+            # Set file output paths
+            scene.render.filepath = render_path
+
+            # Make sure the output directory exists
+            os.makedirs(output_dir, exist_ok=True)
+
+            # Render
+            bpy.ops.render.render(write_still=True)
+
+            context.view_layer.objects.active = empty
+            empty.rotation_euler[2] += math.radians(stepsize)
+
+    # ensure that all objects have materials, if not then add a default
+    # one.
+    textured_mat = create_solid_color_material("default texture", [0.6, 0.6, 0.6, 1])
+
+    for obj in get_scene_meshes():
+        if obj.active_material is None:
+            obj.active_material = textured_mat
+
+    render_views("textured")
+
+
+def enable_gpus(device_type, use_cpus=False):
+    preferences = bpy.context.preferences
+    cycles_preferences = preferences.addons["cycles"].preferences
+    cycles_preferences.refresh_devices()
+    try:
+        devices = cycles_preferences.devices
+    except:
+        print("No devices detected")
+        if device_type == "CPU":
+            return []
+        else:
+            raise RuntimeError(f"No devices detected, set use_cpus to True")
+
+    assert device_type in [
+        "CUDA",
+        "METAL",
+        "OPENCL",
+        "CPU",
+        "NONE",
+    ], f"Unsupported device type: {device_type}"
+
+    try:
+        # print(devices)
+        iter(devices)
+    except TypeError:
+        # print("Single GPU Detected")
+        devices = [devices]
+
+    activated_gpus = []
+    for device in devices:
+        if device.type == "CPU":
+            device.use = use_cpus
+        else:
+            device.use = True
+            activated_gpus.append(device.name)
+
+    if device_type == "CUDA":
+        cycles_preferences.compute_device_type = "CUDA"
+        bpy.context.scene.cycles.device = "GPU"
+    elif device_type == "METAL":
+        cycles_preferences.compute_device_type = "METAL"
+        bpy.context.scene.cycles.device = "GPU"
+    elif device_type == "OPENCL":
+        cycles_preferences.compute_device_type = "OPENCL"
+        bpy.context.scene.cycles.device = "GPU"
+    else:
+        raise RuntimeError(f"Unsupported device type: {device_type}")
+
+    return activated_gpus
+
+
+def set_render_settings(engine, resolution):
+    # Set render settings
+    render.engine = engine  #
+    render.image_settings.file_format = "PNG"
+    render.image_settings.color_mode = "RGBA"
+    render.resolution_x = resolution
+    render.resolution_y = resolution
+    render.resolution_percentage = 100
+
+    # Set cycles settings
+    scene.cycles.device = "GPU"
+    scene.cycles.use_adaptive_sampling = True
+    scene.cycles.adaptive_threshold = 0.1
+    scene.cycles.samples = 64
+    scene.cycles.adaptive_min_samples = 1
+    scene.cycles.filter_width = 2
+    scene.cycles.use_fast_gi = True
+    scene.cycles.fast_gi_method = "REPLACE"
+    world.light_settings.ao_factor = 1.0
+    world.light_settings.distance = 10
+    scene.cycles.use_denoising = True  # ML denoising
+    scene.cycles.denoising_use_gpu = True
+
+    # bake existing frames for faster future renders
+    scene.render.use_persistent_data = True
+
+    # Set eevee settings
+    scene.eevee.use_shadows = True
+    scene.eevee.use_raytracing = True
+    scene.eevee.ray_tracing_options.use_denoise = True
+    scene.eevee.use_fast_gi = True
+    scene.eevee.fast_gi_method = "GLOBAL_ILLUMINATION"
+    scene.eevee.ray_tracing_options.trace_max_roughness = 0.5
+    scene.eevee.fast_gi_resolution = "2"
+    scene.eevee.fast_gi_ray_count = 2
+    scene.eevee.fast_gi_step_count = 8
+
+
+def print_devices():
+    print("Devices:")
+    preferences = bpy.context.preferences
+    cycles_preferences = preferences.addons["cycles"].preferences
+    cycles_preferences.refresh_devices()
+
+    devices = cycles_preferences.devices
+    for device in devices:
+        print(f'   [{device.id}]<{device.type}> "{device.name}" Using: {device.use}')
+
+    print(f"Compute device type: {cycles_preferences.compute_device_type}")
+    print(f"Cycles device: {bpy.context.scene.cycles.device}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--object_path",
+        type=str,
+        required=False,
+        help="Path to the object file",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=True,
+        help="Path to the directory where the rendered images and metadata will be saved.",
+    )
+    parser.add_argument(
+        "--engine",
+        type=str,
+        default="BLENDER_EEVEE_NEXT",  # BLENDER_BLENDER_EEVEE_NEXT rasterization, better than nvdifrast, CYCLES
+        choices=["CYCLES", "BLENDER_EEVEE_NEXT"],
+    )
+    parser.add_argument(
+        "--num_renders",
+        type=int,
+        default=12,
+        help="Number of renders to save of the object.",
+    )
+    parser.add_argument(
+        "--render_resolution",
+        type=int,
+        default=512,
+        help="Resolution of the rendered images.",
+    )
+    parser.add_argument(
+        "--transparent_background",
+        action="store_true",
+        help="Whether to use transparent background",
+    )
+    parser.add_argument(
+        "--environment_map",
+        default=None,
+        type=str,
+        help="Use the given environment map for lighting",
+    )
+
+    argv = sys.argv[sys.argv.index("--") + 1 :]
+    args = parser.parse_args(argv)
+
+    context = bpy.context
+    scene = context.scene
+    render = scene.render
+    world = bpy.data.worlds["World"]
+
+    set_render_settings(args.engine, args.render_resolution)
+
+    # detect platform and activate GPUs
+    platform = platform.system()
+    if platform == "Darwin":
+        activated_gpus = enable_gpus("METAL", use_cpus=True)
+    elif platform == "Linux":
+        activated_gpus = enable_gpus("CUDA", use_cpus=False)
+    else:
+        raise RuntimeError("Unsupported platform")
+    print(f"Activated GPUs: {activated_gpus}")
+
+    print_devices()
+
+    render_object(
+        object_file=args.object_path,
+        num_renders=args.num_renders,
+        output_dir=args.output_dir,
+        transparent_background=args.transparent_background,
+        environment_map=args.environment_map,
+    )

cube/cube3d/renderer/renderer.py

@@ -0,0 +1,88 @@
+import argparse
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+from PIL import Image
+
+
+def render_asset(
+    asset_path,
+    output_dir,
+    nviews=24,
+    img_resolution=512,
+):
+    """
+    Render given asset into output_dir and return the saved image paths.
+    Assumes that blender is installed and is in your path.
+
+        nviews : number of views to render
+        img_resolution : resolution of each rendered view in pixels
+    """
+
+    curr_file_path = __file__
+    curr_dir = os.path.dirname(curr_file_path)
+
+    command = [
+        "blender",
+        "--background",
+        "-noaudio",
+        "--python",
+        f"{curr_dir}/blender_script.py",
+        "--",
+        "--object_path",
+        asset_path,
+        "--num_renders",
+        str(nviews),
+        "--output_dir",
+        output_dir,
+        "--render_resolution",
+        str(img_resolution),
+        "--transparent_background",
+        "--engine",
+        "CYCLES",
+    ]
+
+    subprocess.run(command, check=True)
+
+    # return the saved images paths
+    images = []
+
+    for i in range(nviews):
+        fp = os.path.abspath(os.path.join(output_dir, f"{i:03d}_textured.png"))
+        images.append(fp)
+
+    return images
+
+
+def save_gif(image_paths, outfile):
+    images = [Image.open(img) for img in image_paths]
+    if len(images) > 1:
+        background = Image.new("RGBA", images[0].size, (255, 255, 255))
+        images = [
+            Image.alpha_composite(background, png).convert("RGB") for png in images
+        ]
+        images[0].save(
+            outfile, save_all=True, append_images=images[1:], duration=100, loop=0
+        )
+
+
+def render_turntable(obj_path, output_dir, output_name="turntable"):
+    """
+    Render a turntable gif of the mesh. Assumes that blender is installed and is in your path.
+        obj_path : path to the obj file
+        output_dir : directory to save the gif. Final image will be saved as `turntable.gif`
+    """
+    image_paths = render_asset(obj_path, output_dir)
+    gif_turntable_outfile = Path(output_dir) / f"{output_name}.gif"
+    save_gif(image_paths, gif_turntable_outfile)
+    return gif_turntable_outfile
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input")
+    parser.add_argument("-o", "--output_dir")
+    args = parser.parse_args(sys.argv[1:])
+    render_turntable(args.input, args.output_dir)

cube/cube3d/vq_vae_encode_decode.py

@@ -0,0 +1,150 @@
+import argparse
+import logging
+
+import numpy as np
+import torch
+import trimesh
+
+from cube3d.inference.utils import load_config, load_model_weights, parse_structured
+from cube3d.model.autoencoder.one_d_autoencoder import OneDAutoEncoder
+
+MESH_SCALE = 0.96
+
+
+def rescale(vertices: np.ndarray, mesh_scale: float = MESH_SCALE) -> np.ndarray:
+    """Rescale the vertices to a cube, e.g., [-1, -1, -1] to [1, 1, 1] when mesh_scale=1.0"""
+    vertices = vertices
+    bbmin = vertices.min(0)
+    bbmax = vertices.max(0)
+    center = (bbmin + bbmax) * 0.5
+    scale = 2.0 * mesh_scale / (bbmax - bbmin).max()
+    vertices = (vertices - center) * scale
+    return vertices
+
+
+def load_scaled_mesh(file_path: str) -> trimesh.Trimesh:
+    """
+    Load a mesh and scale it to a unit cube, and clean the mesh.
+    Parameters:
+        file_obj: str | IO
+        file_type: str
+    Returns:
+        mesh: trimesh.Trimesh
+    """
+    mesh: trimesh.Trimesh = trimesh.load(file_path, force="mesh")
+    mesh.remove_infinite_values()
+    mesh.update_faces(mesh.nondegenerate_faces())
+    mesh.update_faces(mesh.unique_faces())
+    mesh.remove_unreferenced_vertices()
+    if len(mesh.vertices) == 0 or len(mesh.faces) == 0:
+        raise ValueError("Mesh has no vertices or faces after cleaning")
+    mesh.vertices = rescale(mesh.vertices)
+    return mesh
+
+
+def load_and_process_mesh(file_path: str, n_samples: int = 8192):
+    """
+    Loads a 3D mesh from the specified file path, samples points from its surface,
+    and processes the sampled points into a point cloud with normals.
+    Args:
+        file_path (str): The file path to the 3D mesh file.
+        n_samples (int, optional): The number of points to sample from the mesh surface. Defaults to 8192.
+    Returns:
+        torch.Tensor: A tensor of shape (1, n_samples, 6) containing the processed point cloud.
+                        Each point consists of its 3D position (x, y, z) and its normal vector (nx, ny, nz).
+    """
+
+    mesh = load_scaled_mesh(file_path)
+    positions, face_indices = trimesh.sample.sample_surface(mesh, n_samples)
+    normals = mesh.face_normals[face_indices]
+    point_cloud = np.concatenate(
+        [positions, normals], axis=1
+    )  # Shape: (num_samples, 6)
+    point_cloud = torch.from_numpy(point_cloud.reshape(1, -1, 6)).float()
+    return point_cloud
+
+
+@torch.inference_mode()
+def run_shape_decode(
+    shape_model: OneDAutoEncoder,
+    output_ids: torch.Tensor,
+    resolution_base: float = 8.0,
+    chunk_size: int = 100_000,
+):
+    """
+    Decodes the shape from the given output IDs and extracts the geometry.
+    Args:
+        shape_model (OneDAutoEncoder): The shape model.
+        output_ids (torch.Tensor): The tensor containing the output IDs.
+        resolution_base (float, optional): The base resolution for geometry extraction. Defaults to 8.43.
+        chunk_size (int, optional): The chunk size for processing. Defaults to 100,000.
+    Returns:
+        tuple: A tuple containing the vertices and faces of the mesh.
+    """
+    shape_ids = (
+        output_ids[:, : shape_model.cfg.num_encoder_latents, ...]
+        .clamp_(0, shape_model.cfg.num_codes - 1)
+        .view(-1, shape_model.cfg.num_encoder_latents)
+    )
+    latents = shape_model.decode_indices(shape_ids)
+    mesh_v_f, _ = shape_model.extract_geometry(
+        latents,
+        resolution_base=resolution_base,
+        chunk_size=chunk_size,
+        use_warp=True,
+    )
+    return mesh_v_f
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="cube shape encode and decode example script"
+    )
+    parser.add_argument(
+        "--mesh-path",
+        type=str,
+        required=True,
+        help="Path to the input mesh file.",
+    )
+    parser.add_argument(
+        "--config-path",
+        type=str,
+        default="cube3d/configs/open_model.yaml",
+        help="Path to the configuration YAML file.",
+    )
+    parser.add_argument(
+        "--shape-ckpt-path",
+        type=str,
+        required=True,
+        help="Path to the shape encoder/decoder checkpoint file.",
+    )
+    parser.add_argument(
+        "--recovered-mesh-path",
+        type=str,
+        default="recovered_mesh.obj",
+        help="Path to save the recovered mesh file.",
+    )
+    args = parser.parse_args()
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    logging.info(f"Using device: {device}")
+
+    cfg = load_config(args.config_path)
+
+    shape_model = OneDAutoEncoder(
+        parse_structured(OneDAutoEncoder.Config, cfg.shape_model)
+    )
+    load_model_weights(
+        shape_model,
+        args.shape_ckpt_path,
+    )
+    shape_model = shape_model.eval().to(device)
+    point_cloud = load_and_process_mesh(args.mesh_path)
+    output = shape_model.encode(point_cloud.to(device))
+    indices = output[3]["indices"]
+    print("Got the following shape indices:")
+    print(indices)
+    print("Indices shape: ", indices.shape)
+    mesh_v_f = run_shape_decode(shape_model, indices)
+    vertices, faces = mesh_v_f[0][0], mesh_v_f[0][1]
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+    mesh.export(args.recovered_mesh_path)

cube/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "cube"
+version = "0.1"
+requires-python = ">=3.7"
+description = "A generative 3D model to accelerate the creation of 3D assets, accessories, and experiences."
+authors = [
+    { name = "Foundation AI", email = "[email protected]" }
+]
+keywords = ["cube"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3.10",
+]
+dependencies = [
+    "numpy",
+    "torch>=2.2.2",
+    "tqdm",
+    "transformers",
+    "omegaconf",
+    "warp-lang",
+    "accelerate>=0.26.0",
+    "scikit-image",
+    "huggingface_hub[cli]",
+    "trimesh"
+]
+[project.optional-dependencies]
+meshlab = ["pymeshlab"]
+lint = ["ruff==0.9.10"]
+
+[tool.setuptools.packages.find]
+where = ["cube3d"]
+include = ["cube/*"]
+namespaces = false

cube/setup.py

@@ -0,0 +1,6 @@
+from setuptools import setup
+
+setup(
+    name="cube",
+    version="0.0.1",
+)